boocode/apps/coder/src/services/provider-snapshot.ts

/**
 * Provider snapshot cache — cold ACP probe per provider + static manifest merge.
 */
import { homedir } from 'node:os';
import type { FastifyBaseLogger } from 'fastify';
import type { Sql } from '../db.js';
import type { Config } from '../config.js';
import {
  getManifestDefaultModeId,
  getManifestModes,
  PROVIDER_MANIFEST,
} from './provider-manifest.js';
import { probeAcpProvider } from './acp-probe.js';
import type { ProviderModel, ProviderSnapshotEntry, AgentCommand } from './provider-types.js';
import { getManifestCommands, mergeCommands } from './provider-commands.js';
import { readQwenSettingsModels } from './qwen-settings.js';
import { getResolvedRegistry, type ResolvedProviderDef } from './provider-config-registry.js';
import { isCommandAvailable } from './command-availability.js';
import { discoverClaudeCommands } from './claude-command-discovery.js';

interface AgentRow {
  name: string;
  install_path: string | null;
  supports_acp: boolean;
  models: ProviderModel[] | null;
  commands: AgentCommand[] | null;
  label: string | null;
  transport: string | null;
  last_probed_at: string | Date | null;
}

export async function fetchLlamaSwapModels(config: Config): Promise<ProviderModel[]> {
  try {
    const res = await fetch(`${config.LLAMA_SWAP_URL}/v1/models`);
    if (!res.ok) return [];
    const parsed = (await res.json()) as { data?: Array<{ id: string }> };
    return (parsed.data ?? []).map((m) => ({ id: m.id, label: m.id }));
  } catch {
    return [];
  }
}

/** Prefix llama-swap model ids so they don't collide with provider-native models. */
export function prefixLlamaSwapModels(models: ProviderModel[]): ProviderModel[] {
  return models.map((m) => ({
    ...m,
    id: m.id.startsWith('llama-swap/') ? m.id : `llama-swap/${m.id}`,
  }));
}

function attachClaudeThinking(models: ProviderModel[]): ProviderModel[] {
  const thinking = PROVIDER_MANIFEST.claude?.thinkingOptions;
  if (!thinking?.length) return models;
  return models.map((m) => ({
    ...m,
    thinkingOptions: thinking,
    defaultThinkingOptionId: 'medium',
  }));
}

export function mergeModels(...lists: ProviderModel[][]): ProviderModel[] {
  const seen = new Set<string>();
  const out: ProviderModel[] = [];
  for (const list of lists) {
    for (const m of list) {
      if (seen.has(m.id)) continue;
      seen.add(m.id);
      out.push(m);
    }
  }
  return out;
}

async function buildProviderEntry(
  resolved: ResolvedProviderDef,
  agentRow: AgentRow | undefined,
  llamaModels: ProviderModel[],
  cwd: string,
  ttlMs: number,
  force: boolean,
): Promise<ProviderSnapshotEntry> {
  const name = resolved.id;
  const isNative = resolved.transport === 'native';
  const fallbackModes = getManifestModes(name);
  const defaultModeId = getManifestDefaultModeId(name);
  const manifestCommands = getManifestCommands(name);
  // Manifest + persisted live ACP commands (captured on a prior cold probe), so
  // the agent's discovered commands show even when the tier-2 probe is skipped.
  const dbCommands = mergeCommands(manifestCommands, agentRow?.commands ?? []);
  const label = agentRow?.label ?? resolved.configLabel ?? resolved.label;
  const descr = resolved.configDescription ? { description: resolved.configDescription } : {};

  // v2.3: config `models` REPLACES the discovered/static list; `additionalModels`
  // MERGES on top. Applied to every ready/installed model list below.
  const withConfigModels = (m: ProviderModel[]): ProviderModel[] => {
    let out = resolved.configModels && resolved.configModels.length > 0 ? resolved.configModels : m;
    if (resolved.configAdditionalModels && resolved.configAdditionalModels.length > 0) {
      out = mergeModels(out, resolved.configAdditionalModels);
    }
    return out;
  };

  // ACP built-ins fall back to PTY transport when the installed binary lacks ACP.
  let transport = resolved.transport;
  if (agentRow && resolved.transport === 'acp' && !agentRow.supports_acp) {
    transport = 'pty';
  }

  // 1. Disabled → unavailable, no probe.
  if (!resolved.enabled) {
    return {
      name, label, ...descr, transport, status: 'unavailable',
      enabled: false, installed: false, models: [], modes: fallbackModes,
      defaultModeId, commands: manifestCommands,
    };
  }

  // 2. Native boocode → always ready (llama-swap models).
  if (isNative) {
    return {
      name, label: resolved.label, transport, status: 'ready',
      enabled: true, installed: true, models: withConfigModels(llamaModels), modes: [],
      defaultModeId: null, commands: manifestCommands,
    };
  }

  // 3. Tier-1 fast availability: installed iff a probed install_path exists or
  // the launch binary is on PATH. No spawn beyond a `which` for custom entries.
  const fast =
    agentRow?.install_path != null ||
    (resolved.launchCommand ? await isCommandAvailable(resolved.launchCommand[0]) : false);

  if (!fast) {
    return {
      name, label, ...descr, transport, status: 'unavailable',
      enabled: true, installed: false, models: [], modes: fallbackModes,
      defaultModeId, commands: manifestCommands,
    };
  }

  // Baseline model precedence (used by claude + non-probe fallbacks).
  let models: ProviderModel[] = [];
  if (resolved.modelSource === 'llama-swap' && resolved.mergeLlamaSwap) {
    models = llamaModels;
  } else if (agentRow?.models?.length) {
    models = agentRow.models;
  } else if (resolved.staticModels) {
    models = resolved.staticModels.map((m) => ({ id: m.id, label: m.label }));
  }

  // claude: static models + thinking options, no ACP probe (unchanged from v2.2).
  if (name === 'claude') {
    // claude is PTY (no ACP discovery) — read its enabled commands + plugin
    // skills from disk live (the snapshot cache rate-limits the fs reads).
    return {
      name, label, transport, status: 'ready', enabled: true, installed: true,
      models: attachClaudeThinking(withConfigModels(models)), modes: fallbackModes, defaultModeId,
      commands: mergeCommands(manifestCommands, discoverClaudeCommands()),
    };
  }

  const canProbeAcp =
    transport === 'acp' &&
    ((agentRow?.install_path != null && agentRow.supports_acp) ||
      (resolved.isCustomAcp && resolved.launchCommand != null));

  if (canProbeAcp) {
    // Tier-2 gate (§4.3): cold ACP probe only on force, staleness, or empty DB
    // models. Otherwise serve DB models + manifest modes/commands — no spawn.
    const lastProbedMs =
      agentRow?.last_probed_at != null ? new Date(agentRow.last_probed_at).getTime() : NaN;
    const stale = Number.isNaN(lastProbedMs) || Date.now() - lastProbedMs > ttlMs;
    const dbEmpty = !(agentRow?.models && agentRow.models.length > 0);
    const runTier2 = force || stale || dbEmpty;

    if (!runTier2) {
      let skipModels = agentRow?.models ?? [];
      if (resolved.mergeLlamaSwap && resolved.modelSource !== 'llama-swap') {
        skipModels = mergeModels(skipModels, prefixLlamaSwapModels(llamaModels));
      } else if (resolved.modelSource === 'llama-swap' && skipModels.length === 0) {
        skipModels = llamaModels;
      }
      return {
        name, label, transport, status: 'ready', enabled: true, installed: true,
        models: withConfigModels(skipModels), modes: fallbackModes, defaultModeId, commands: dbCommands,
      };
    }

    const probeTarget =
      resolved.isCustomAcp && resolved.launchCommand
        ? resolved.launchCommand[0]
        : agentRow!.install_path!;
    const probe = await probeAcpProvider(name, probeTarget, cwd);

    let probeModels = probe.models.length > 0 ? probe.models : models;
    if (name === 'qwen') {
      probeModels = mergeModels(probeModels, await readQwenSettingsModels());
    }
    if (resolved.mergeLlamaSwap && resolved.modelSource !== 'llama-swap') {
      const nativeModels = probe.models.length > 0 ? probe.models : probeModels;
      probeModels = mergeModels(nativeModels, prefixLlamaSwapModels(llamaModels));
    }

    return {
      name, label, transport,
      status: probe.ok ? 'ready' : 'error',
      enabled: true, installed: true,
      models: withConfigModels(probeModels),
      modes: probe.modes.length > 0 ? probe.modes : fallbackModes,
      defaultModeId: probe.defaultModeId ?? defaultModeId,
      commands: mergeCommands(manifestCommands, probe.commands),
      ...(probe.error ? { error: probe.error } : {}),
      fetchedAt: new Date().toISOString(),
    };
  }

  // PTY-only fallback (e.g. qwen without ACP) — installed + ready.
  if (name === 'qwen' && models.length === 0) {
    models = await readQwenSettingsModels();
  }
  return {
    name, label, transport, status: 'ready', enabled: true, installed: true,
    models: withConfigModels(models), modes: fallbackModes, defaultModeId, commands: dbCommands,
  };
}

const snapshotCache = new Map<string, { at: number; entries: ProviderSnapshotEntry[] }>();
const snapshotInflight = new Map<string, Promise<ProviderSnapshotEntry[]>>();
const CACHE_TTL_MS = 5 * 60_000;

export async function getProviderSnapshot(
  sql: Sql,
  config: Config,
  cwd?: string,
  force = false,
): Promise<ProviderSnapshotEntry[]> {
  const resolvedCwd = cwd?.trim() || homedir();
  const cacheKey = resolvedCwd;
  const cached = snapshotCache.get(cacheKey);
  if (!force && cached && Date.now() - cached.at < CACHE_TTL_MS) {
    return cached.entries;
  }

  const inflight = snapshotInflight.get(cacheKey);
  if (!force && inflight) {
    return inflight;
  }

  const build = async (): Promise<ProviderSnapshotEntry[]> => {
    const llamaModels = await fetchLlamaSwapModels(config);
    const agents = await sql<AgentRow[]>`
      SELECT name, install_path, supports_acp, models, commands, label, transport, last_probed_at FROM available_agents
    `;
    const agentMap = new Map(agents.map((a) => [a.name, a]));
    const ttlMs = config.PROVIDER_PROBE_TTL_MS;

    const entries = await Promise.all(
      [...getResolvedRegistry().values()].map((resolved) =>
        buildProviderEntry(resolved, agentMap.get(resolved.id), llamaModels, resolvedCwd, ttlMs, force),
      ),
    );

    snapshotCache.set(cacheKey, { at: Date.now(), entries });
    return entries;
  };

  const promise = build().finally(() => {
    snapshotInflight.delete(cacheKey);
  });
  snapshotInflight.set(cacheKey, promise);

  // Await the build (force or cache-miss) and return terminal entries. The sync
  // `loading` return (design §4.4) is DEFERRED until Phase 5 ships the client
  // poll that resolves it: without that poll, a single fetch lands on
  // installed:false `loading` entries, which AgentComposerBar filters out
  // (`e.installed && ...`) → empty picker. Builds stay fast via the tier-2 skip
  // once available_agents.models is warm.
  return promise;
}

export function clearProviderSnapshotCache(): void {
  snapshotCache.clear();
  snapshotInflight.clear();
}

/**
 * Read-only peek into the warm snapshot cache for one provider (no build, no
 * probe). Used by the diagnostic route to report the last computed probe error
 * without spawning anything. Returns undefined on a cold cache / unknown name.
 */
export function peekSnapshotEntry(name: string, cwd?: string): ProviderSnapshotEntry | undefined {
  const resolvedCwd = cwd?.trim() || homedir();
  return snapshotCache.get(resolvedCwd)?.entries.find((e) => e.name === name);
}

/** Persist probed model lists back to available_agents for fast legacy reads. */
export async function persistProbedModels(
  sql: Sql,
  entries: ProviderSnapshotEntry[],
  log: FastifyBaseLogger,
): Promise<void> {
  let count = 0;
  for (const entry of entries) {
    if (entry.name === 'boocode') continue;
    let persisted = false;
    if (entry.models.length > 0) {
      const flatModels = entry.models.map(({ id, label }) => ({ id, label }));
      await sql`
        UPDATE available_agents
        SET models = ${sql.json(flatModels as never)}, last_probed_at = clock_timestamp()
        WHERE name = ${entry.name}
      `;
      persisted = true;
    }
    // Persist captured ACP commands so they survive the tier-2 probe skip and
    // show without a dispatch. Only when non-empty — never clobber a prior set.
    if (entry.commands.length > 0) {
      const flatCommands = entry.commands.map((c) => ({
        name: c.name,
        ...(c.description ? { description: c.description } : {}),
      }));
      await sql`
        UPDATE available_agents
        SET commands = ${sql.json(flatCommands as never)}, last_probed_at = clock_timestamp()
        WHERE name = ${entry.name}
      `;
      persisted = true;
    }
    if (persisted) count++;
  }
  if (count > 0) {
    log.info({ count }, 'provider-snapshot: persisted models/commands to available_agents');
  }
}