Codewalkers/apps/server/agent/credential-handler.ts

/**
 * CredentialHandler — Account selection, credential management, and exhaustion handling.
 *
 * Extracted from MultiProviderAgentManager. Handles account lifecycle:
 * selecting the next available account, writing credentials to disk,
 * ensuring they're fresh, and marking accounts as exhausted on failure.
 */

import { readFile } from 'node:fs/promises';
import { existsSync } from 'node:fs';
import { join } from 'node:path';
import type { AccountRepository } from '../db/repositories/account-repository.js';
import type { AccountCredentialManager } from './credentials/types.js';
import type { Account } from '../db/schema.js';
import { ensureAccountCredentials } from './accounts/usage.js';
import { getAccountConfigDir } from './accounts/paths.js';
import { setupAccountConfigDir } from './accounts/setup.js';
import { createModuleLogger } from '../logger/index.js';

const log = createModuleLogger('credential-handler');

/** Default exhaustion duration: 5 hours */
const DEFAULT_EXHAUSTION_HOURS = 5;

export class CredentialHandler {
  constructor(
    private workspaceRoot: string,
    private accountRepository?: AccountRepository,
    private credentialManager?: AccountCredentialManager,
  ) {}

  /**
   * Select the next available account for a provider.
   * Clears expired exhaustion, returns least-recently-used non-exhausted account.
   * Returns null if no accounts are available.
   */
  async selectAccount(providerName: string): Promise<{ account: Account; accountId: string; configDir: string } | null> {
    if (!this.accountRepository) return null;

    await this.accountRepository.clearExpiredExhaustion();
    const account = await this.accountRepository.findNextAvailable(providerName);
    if (!account) return null;

    const configDir = getAccountConfigDir(this.workspaceRoot, account.id);
    await this.accountRepository.updateLastUsed(account.id);

    return { account, accountId: account.id, configDir };
  }

  /**
   * Write account credentials from DB to the convention-based config directory.
   * Must be called before ensureCredentials so the files exist on disk.
   */
  async writeCredentialsToDisk(account: Account, configDir: string): Promise<void> {
    if (account.configJson && account.credentials) {
      await setupAccountConfigDir(configDir, {
        configJson: JSON.parse(account.configJson),
        credentials: account.credentials,
      });
      log.debug({ accountId: account.id, configDir }, 'wrote account credentials from DB to disk');
    } else {
      log.warn({ accountId: account.id }, 'account has no stored credentials in DB');
    }
  }

  /**
   * Read refreshed credentials from disk and persist back to DB.
   * Called after credential refresh to keep DB in sync.
   */
  async persistRefreshedCredentials(accountId: string, configDir: string): Promise<void> {
    if (!this.accountRepository) return;
    try {
      const credPath = join(configDir, '.credentials.json');
      const credentials = await readFile(credPath, 'utf-8');
      await this.accountRepository.updateCredentials(accountId, credentials);
      log.debug({ accountId }, 'persisted refreshed credentials back to DB');
    } catch (err) {
      log.warn({ accountId, err: err instanceof Error ? err.message : String(err) }, 'failed to persist refreshed credentials to DB');
    }
  }

  /**
   * Ensure credentials are valid before spawn/resume.
   * Uses credentialManager if available, otherwise falls back to legacy function.
   * Returns { valid, refreshed } so callers can persist refresh back to DB.
   */
  async ensureCredentials(configDir: string, accountId?: string): Promise<{ valid: boolean; refreshed: boolean }> {
    if (this.credentialManager) {
      const result = await this.credentialManager.ensureValid(configDir, accountId);
      return { valid: result.valid, refreshed: result.refreshed };
    }
    const valid = await ensureAccountCredentials(configDir);
    return { valid, refreshed: false };
  }

  /**
   * Read the access token from a config directory's .credentials.json.
   * Returns null if credentials file is missing or malformed.
   * Used for CLAUDE_CODE_OAUTH_TOKEN env var injection.
   */
  async readAccessToken(configDir: string): Promise<string | null> {
    try {
      const credPath = join(configDir, '.credentials.json');
      if (!existsSync(credPath)) return null;
      const raw = await readFile(credPath, 'utf-8');
      const parsed = JSON.parse(raw);
      return parsed.claudeAiOauth?.accessToken ?? null;
    } catch {
      return null;
    }
  }

  /**
   * Prepare process environment with account credentials.
   * Writes credentials to disk, ensures freshness, injects OAuth token.
   * Used by spawn, resumeForCommit, and resumeInternal.
   */
  async prepareProcessEnv(
    providerEnv: Record<string, string>,
    provider: { configDirEnv?: string },
    accountId: string | null,
  ): Promise<{ processEnv: Record<string, string>; accountConfigDir: string | null }> {
    const processEnv: Record<string, string> = { ...providerEnv };
    let accountConfigDir: string | null = null;

    if (accountId && provider.configDirEnv && this.accountRepository) {
      accountConfigDir = getAccountConfigDir(this.workspaceRoot, accountId);
      const account = await this.accountRepository.findById(accountId);
      if (account) {
        await this.writeCredentialsToDisk(account, accountConfigDir);
      }
      processEnv[provider.configDirEnv] = accountConfigDir;

      const { valid, refreshed } = await this.ensureCredentials(accountConfigDir, accountId);
      if (!valid) {
        log.warn({ accountId }, 'failed to refresh credentials');
      }
      if (refreshed) {
        await this.persistRefreshedCredentials(accountId, accountConfigDir);
      }

      const accessToken = await this.readAccessToken(accountConfigDir);
      if (accessToken) {
        processEnv['CLAUDE_CODE_OAUTH_TOKEN'] = accessToken;
        log.debug({ accountId }, 'CLAUDE_CODE_OAUTH_TOKEN injected');
      }
    }

    return { processEnv, accountConfigDir };
  }

  /**
   * Check if an error message indicates usage limit exhaustion.
   */
  isUsageLimitError(errorMessage: string): boolean {
    const patterns = [
      'usage limit',
      'rate limit',
      'quota exceeded',
      'too many requests',
      'capacity',
      'exhausted',
    ];
    const lower = errorMessage.toLowerCase();
    return patterns.some((p) => lower.includes(p));
  }

  /**
   * Handle account exhaustion: mark current account exhausted and find next available.
   * Returns the new account info if failover succeeded, null otherwise.
   * Does NOT re-spawn — the caller (manager) handles that.
   */
  async handleExhaustion(
    accountId: string,
    providerName: string,
  ): Promise<{ account: Account; accountId: string; configDir: string } | null> {
    if (!this.accountRepository) return null;

    log.warn({ accountId, provider: providerName }, 'account exhausted, attempting failover');

    // Mark current account as exhausted
    const exhaustedUntil = new Date(Date.now() + DEFAULT_EXHAUSTION_HOURS * 60 * 60 * 1000);
    await this.accountRepository.markExhausted(accountId, exhaustedUntil);

    // Find next available account
    const nextAccount = await this.accountRepository.findNextAvailable(providerName);
    if (!nextAccount) {
      log.warn({ accountId }, 'account failover failed, no accounts available');
      return null;
    }
    log.info({ previousAccountId: accountId, newAccountId: nextAccount.id }, 'account failover successful');

    // Write credentials and ensure they're fresh
    const nextConfigDir = getAccountConfigDir(this.workspaceRoot, nextAccount.id);
    await this.writeCredentialsToDisk(nextAccount, nextConfigDir);
    const { valid, refreshed } = await this.ensureCredentials(nextConfigDir, nextAccount.id);
    if (!valid) {
      log.warn({ newAccountId: nextAccount.id }, 'failed to refresh failover account credentials');
      return null;
    }
    if (refreshed) {
      await this.persistRefreshedCredentials(nextAccount.id, nextConfigDir);
    }

    await this.accountRepository.updateLastUsed(nextAccount.id);

    return { account: nextAccount, accountId: nextAccount.id, configDir: nextConfigDir };
  }
}