okengine 0.18.4 → 0.19.0

package/src/cli/ai-setup/apply.ts

@@ -4,11 +4,17 @@
 
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { dirname, join } from "node:path";
-import { OLLAMA_IMAGE } from "../../docker/recipes/index.ts";
+import { getAiProviderEntry } from "../../elements/ai/providers.ts";
+import { CLOUD_PROVIDERS } from "./catalog.ts";
 
 /** Choices applied to the project. */
 export type AiSetupApplyInput = {
-  readonly driver: "ollama" | "anthropic" | "openai-compatible" | "mock";
+  readonly driver: "anthropic" | "openai-compatible" | "mock";
+  /**
+   * Name written to `ai.model({ provider })` — registry id (`openrouter`,
+   * `openai`, …), or `openai-compatible` / `local` for self-hosted.
+   */
+  readonly provider?: string;
   readonly baseUrl?: string;
   readonly chatModel?: string;
   readonly visionModel?: string | null;
@@ -16,7 +22,7 @@ export type AiSetupApplyInput = {
   readonly apiKeyEnv?: string;
   /** When set with {@link apiKeyEnv}, writes the token into `.env.local`. */
   readonly apiKey?: string;
-  /** Optional `images.ai` pin (llama.cpp / Ollama / vLLM / SGLang). */
+  /** Optional `images.ai` pin (legacy; cloud / custom / lmstudio leave unset). */
   readonly image?: string;
 };
 
@@ -56,8 +62,9 @@ export function applyAiSetup(
     config = upsertAiDrivers(config, input.driver);
     if (input.image) {
       config = upsertImage(config, "ai", input.image);
-    } else if (input.driver === "ollama") {
-      config = upsertImage(config, "ai", OLLAMA_IMAGE);
+    } else {
+      // Cloud / host-side providers must not leave a leftover `images.ai` pin.
+      config = removeImage(config, "ai");
     }
     writeFileSync(configPath, config, "utf8");
   }
@@ -92,6 +99,9 @@ export function applyAiSetup(
   writeFileSync(envPath, env.endsWith("\n") ? env : `${env}\n`, "utf8");
 
   const aiTsPath = writeAiModels(cwd, input);
+  if (input.apiKeyEnv) {
+    ensureAiApiKeyVaultSecret(cwd, input.apiKeyEnv);
+  }
 
   return { configPath, envPath, aiTsPath };
 }
@@ -143,6 +153,19 @@ export function upsertImage(source: string, key: string, image: string): string
   return source.replace(imagesRe, `images: {${body}\n${line}\n  }`);
 }
 
+/**
+ * Drop one flat image role pin (e.g. clear `images.ai` for cloud AI).
+ *
+ * @param source - Config source
+ * @param key - Image role (`ai`, `pgdog`, … — not nested `store.*`)
+ */
+export function removeImage(source: string, key: string): string {
+  const keyLit = key.includes(".") ? `"${key}"` : key;
+  const lineRe = new RegExp(`\\n[ \\t]*${keyLit}\\s*:\\s*"[^"]*",?`);
+  if (!lineRe.test(source)) return source;
+  return source.replace(lineRe, "");
+}
+
 /** Options for {@link upsertEnv}. */
 export type UpsertEnvOptions = {
   /**
@@ -181,34 +204,40 @@ export function upsertEnv(
  */
 export function renderAiTs(input: AiSetupApplyInput): string {
   const chat = input.chatModel ?? "default";
-  const provider =
-    input.driver === "ollama"
-      ? "ollama"
-      : input.driver === "anthropic"
-        ? "anthropic"
-        : input.driver === "mock"
-          ? "mock"
-          : "openai-compatible";
+  const provider = resolveDeclProvider(input);
+  const apiKeyEnv = input.apiKeyEnv;
+  const registryCloud = isRegistryCloudProvider(provider);
+  const nativeAnthropic = input.driver === "anthropic";
+  const localPrimary = isSelfHostedLocal(input);
 
-  const lines = [
-    `import { ai } from "okengine";`,
-    ``,
-    `/** Cloud OpenAI-compatible binding (OpenAI / Groq / OpenRouter / …). */`,
+  const smartLines = [
     `export const smart = ai.model("smart", {`,
     `  provider: "${provider}",`,
+    ...(nativeAnthropic ? [`  driverId: "anthropic",`] : []),
     `  model: process.env.OKE_AI_CLOUD_MODEL ?? process.env.OKE_AI_MODEL ?? "${chat}",`,
-    `  baseUrl: process.env.OPENAI_BASE_URL?.trim() || "https://api.openai.com/v1",`,
-    `  ...(process.env.OPENAI_API_KEY?.trim()`,
-    `    ? { apiKey: process.env.OPENAI_API_KEY.trim() }`,
-    `    : {}),`,
+    ...smartBaseUrlLines(input, { registryCloud, localPrimary }),
+    ...apiKeySpreadLines(apiKeyEnv),
     `});`,
-    ``,
-    `/** Local inference binding (docker llama.cpp / Ollama via \`OKE_AI_URL\`). */`,
+  ];
+
+  const localProvider = "openai-compatible";
+  const localDefaultModel = localPrimary ? chat : "local-model";
+  const localLines = [
     `export const local = ai.model("local", {`,
-    `  provider: "${provider === "ollama" ? "ollama" : "openai-compatible"}",`,
-    `  model: process.env.OKE_AI_LOCAL_MODEL ?? "${chat}",`,
+    `  provider: "${localProvider}",`,
+    `  model: process.env.OKE_AI_LOCAL_MODEL ?? "${localDefaultModel}",`,
     `  ...(process.env.OKE_AI_URL?.trim() ? { baseUrl: process.env.OKE_AI_URL.trim() } : {}),`,
     `});`,
+  ];
+
+  const lines = [
+    `import { ai } from "okengine";`,
+    ``,
+    `/** Primary model binding (${provider}). */`,
+    ...smartLines,
+    ``,
+    `/** Local OpenAI-compatible binding (via \`OKE_AI_URL\`). */`,
+    ...localLines,
     ``,
     `/** Advanced Notes summarize — used by \`notes.summarize\` via \`fx.ask\`. */`,
     `export const summarizeNote = smart.prompt("summarize-note", {`,
@@ -222,7 +251,9 @@ export function renderAiTs(input: AiSetupApplyInput): string {
       ``,
       `export const vision = ai.model("vision", {`,
       `  provider: "${provider}",`,
+      ...(nativeAnthropic ? [`  driverId: "anthropic",`] : []),
       `  model: process.env.OKE_AI_VISION_MODEL ?? "${input.visionModel}",`,
+      ...apiKeySpreadLines(apiKeyEnv),
       `});`,
     );
   }
@@ -233,6 +264,7 @@ export function renderAiTs(input: AiSetupApplyInput): string {
       `export const embedModel = ai.model("embed", {`,
       `  provider: "${provider}",`,
       `  model: process.env.OKE_AI_EMBED_MODEL ?? "${input.embedModel}",`,
+      ...apiKeySpreadLines(apiKeyEnv),
       `});`,
       ``,
       `export const docsEmbed = ai.embed("docs", { model: embedModel });`,
@@ -243,6 +275,163 @@ export function renderAiTs(input: AiSetupApplyInput): string {
   return `${lines.join("\n")}\n`;
 }
 
+/**
+ * Provider string for generated `ai.model` declarations.
+ *
+ * @param input - Setup choices
+ */
+export function resolveDeclProvider(input: AiSetupApplyInput): string {
+  if (input.provider !== undefined && input.provider.length > 0) return input.provider;
+  if (input.driver === "anthropic") return "anthropic";
+  if (input.driver === "mock") return "mock";
+  return "openai-compatible";
+}
+
+function isRegistryCloudProvider(provider: string): boolean {
+  return getAiProviderEntry(provider) !== undefined;
+}
+
+function isSelfHostedLocal(input: AiSetupApplyInput): boolean {
+  return input.image !== undefined;
+}
+
+function smartBaseUrlLines(
+  input: AiSetupApplyInput,
+  flags: { readonly registryCloud: boolean; readonly localPrimary: boolean },
+): string[] {
+  if (input.driver === "anthropic") {
+    return [];
+  }
+  if (flags.registryCloud) {
+    // Registry auto-resolves; optional env override for proxies/mirrors.
+    return [
+      `  ...(process.env.OPENAI_BASE_URL?.trim() ? { baseUrl: process.env.OPENAI_BASE_URL.trim() } : {}),`,
+    ];
+  }
+  if (flags.localPrimary) {
+    return [
+      `  ...(process.env.OKE_AI_URL?.trim() ? { baseUrl: process.env.OKE_AI_URL.trim() } : {}),`,
+    ];
+  }
+  // custom / lmstudio — prefer OPENAI_BASE_URL, fall back to setup-time URL when present
+  if (input.baseUrl) {
+    return [`  baseUrl: process.env.OPENAI_BASE_URL?.trim() || ${JSON.stringify(input.baseUrl)},`];
+  }
+  return [
+    `  ...(process.env.OPENAI_BASE_URL?.trim() ? { baseUrl: process.env.OPENAI_BASE_URL.trim() } : {}),`,
+  ];
+}
+
+function apiKeySpreadLines(apiKeyEnv: string | undefined): string[] {
+  if (!apiKeyEnv) return [];
+  return [
+    `  ...(process.env.${apiKeyEnv}?.trim()`,
+    `    ? { apiKey: process.env.${apiKeyEnv}.trim() }`,
+    `    : {}),`,
+  ];
+}
+
+/**
+ * `OPENROUTER_API_KEY` → `openrouterApiKey`.
+ *
+ * @param envName - Env / vault contract name
+ */
+export function envNameToCamelBinding(envName: string): string {
+  const parts = envName
+    .toLowerCase()
+    .split(/_+/g)
+    .filter((p) => p.length > 0);
+  if (parts.length === 0) return "apiKey";
+  return parts
+    .map((part, i) => (i === 0 ? part : `${part[0]!.toUpperCase()}${part.slice(1)}`))
+    .join("");
+}
+
+/**
+ * Human description for a provider API key env.
+ *
+ * @param envName - Env / vault contract name
+ */
+export function apiKeyEnvDescription(envName: string): string {
+  const meta = CLOUD_PROVIDERS.find((p) => p.apiKeyEnv === envName);
+  if (meta) return `${meta.label} API key`;
+  const stem = envName.replace(/_API_KEY$/i, "").replace(/_/g, " ");
+  return `${stem} API key`;
+}
+
+/**
+ * Render a `vault.secret` block for an AI provider key (no `dev:` stub so
+ * missing values become Vault boot gaps on `oke dev`).
+ *
+ * @param apiKeyEnv - Env / contract name
+ */
+export function renderAiApiKeyVaultSecret(apiKeyEnv: string): string {
+  const binding = envNameToCamelBinding(apiKeyEnv);
+  const description = apiKeyEnvDescription(apiKeyEnv);
+  return `/** ${description} — no \`dev:\` stub so first \`oke dev\` asks via Vault gaps. */
+export const ${binding} = vault.secret(${JSON.stringify(apiKeyEnv)}, {
+  description: ${JSON.stringify(description)},
+  rotate: "90d",
+});
+`;
+}
+
+/**
+ * True when `source` already declares `vault.secret("ENV")` (or `vault("ENV")`).
+ *
+ * @param source - TypeScript source
+ * @param apiKeyEnv - Contract name
+ */
+export function hasAiApiKeyVaultSecret(source: string, apiKeyEnv: string): boolean {
+  const lit = apiKeyEnv.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+  return new RegExp(`\\bvault(?:\\.secret)?\\s*\\(\\s*["']${lit}["']`).test(source);
+}
+
+/**
+ * Merge a vault.secret contract for the AI API key into core / vault sources.
+ *
+ * @param existing - Current module source
+ * @param apiKeyEnv - Env / contract name
+ */
+export function mergeAiApiKeyVaultSecret(existing: string, apiKeyEnv: string): string {
+  if (hasAiApiKeyVaultSecret(existing, apiKeyEnv)) return existing;
+  let next = ensureNamedOkengineImport(existing, "vault");
+  const block = renderAiApiKeyVaultSecret(apiKeyEnv);
+  const vaultHeading = /(\/\/ --- Vault[^\n]*\n)/;
+  if (vaultHeading.test(next)) {
+    return next.replace(vaultHeading, `$1\n${block}`);
+  }
+  const aiHeading = /(\/\/ --- AI[^\n]*\n)/;
+  if (aiHeading.test(next)) {
+    return next.replace(aiHeading, `${block}\n$1`);
+  }
+  return `${next.trimEnd()}\n\n${block}`;
+}
+
+/**
+ * Declare `vault.secret(apiKeyEnv)` in the project so missing keys surface as
+ * Vault gaps (and Console Vault) — never silently optional.
+ *
+ * Prefers `src/core/vault.ts` when present, else merges into `src/core.ts`.
+ *
+ * @param cwd - Project root
+ * @param apiKeyEnv - Env / contract name
+ */
+export function ensureAiApiKeyVaultSecret(cwd: string, apiKeyEnv: string): void {
+  const coreVaultPath = join(cwd, "src", "core", "vault.ts");
+  if (existsSync(coreVaultPath)) {
+    const prev = readFileSync(coreVaultPath, "utf8");
+    const next = mergeAiApiKeyVaultSecret(prev, apiKeyEnv);
+    if (next !== prev) writeFileSync(coreVaultPath, next, "utf8");
+    return;
+  }
+  const coreTsPath = join(cwd, "src", "core.ts");
+  if (!existsSync(coreTsPath)) return;
+  const prev = readFileSync(coreTsPath, "utf8");
+  const next = mergeAiApiKeyVaultSecret(prev, apiKeyEnv);
+  if (next !== prev) writeFileSync(coreTsPath, next, "utf8");
+}
+
 /**
  * Write AI model declarations into `src/core/ai.ts` when that split exists
  * (so a thin `src/core.ts` barrel stays a re-export), else `src/core.ts`,
@@ -260,14 +449,7 @@ function writeAiModels(cwd: string, input: AiSetupApplyInput): string {
 
   if (existsSync(coreAiPath)) {
     const existing = readFileSync(coreAiPath, "utf8");
-    if (hasAiModels(existing)) {
-      const withPrompt = ensureSummarizeNotePrompt(existing);
-      if (withPrompt !== existing) {
-        writeFileSync(coreAiPath, withPrompt, "utf8");
-      }
-    } else {
-      writeFileSync(coreAiPath, mergeAiIntoCore(existing, rendered), "utf8");
-    }
+    writeFileSync(coreAiPath, resolveAiCoreSource(existing, rendered), "utf8");
     ensureCoreBarrelExportsAi(cwd);
     ensureCoreImported(cwd);
     return coreAiPath;
@@ -285,14 +467,7 @@ function writeAiModels(cwd: string, input: AiSetupApplyInput): string {
   mkdirSync(dirname(coreTsPath), { recursive: true });
   if (existsSync(coreTsPath)) {
     const existing = readFileSync(coreTsPath, "utf8");
-    if (hasAiModels(existing)) {
-      const withPrompt = ensureSummarizeNotePrompt(existing);
-      if (withPrompt !== existing) {
-        writeFileSync(coreTsPath, withPrompt, "utf8");
-      }
-      return coreTsPath;
-    }
-    writeFileSync(coreTsPath, mergeAiIntoCore(existing, rendered), "utf8");
+    writeFileSync(coreTsPath, resolveAiCoreSource(existing, rendered), "utf8");
   } else {
     writeFileSync(coreTsPath, rendered, "utf8");
   }
@@ -300,17 +475,82 @@ function writeAiModels(cwd: string, input: AiSetupApplyInput): string {
   return coreTsPath;
 }
 
+/**
+ * Decide whether to merge a full AI module, repair a broken stub, or only
+ * backfill `local` / `summarizeNote` on an already-complete core.
+ *
+ * @param existing - Current core / AI sidecar source
+ * @param rendered - Output of {@link renderAiTs}
+ */
+export function resolveAiCoreSource(existing: string, rendered: string): string {
+  if (isIncompleteAiSetup(existing)) {
+    return mergeAiIntoCore(stripIncompleteAiExports(existing), rendered);
+  }
+  if (hasAiModels(existing)) {
+    return ensureSummarizeNotePrompt(existing);
+  }
+  return mergeAiIntoCore(existing, rendered);
+}
+
 /**
  * @param source - Existing TypeScript
  */
 function hasAiModels(source: string): boolean {
+  // Require real exports — template comments like `//   ai.model("smart", …)`
+  // must not look like an already-configured core.
   return (
-    /\bai\.model\s*\(/.test(source) ||
+    /\bexport\s+const\s+(?:smart|local|vision|embedModel)\s*=\s*ai\.model\s*\(/.test(source) ||
     /from\s+["']\.\/(?:core\/)?ai["']/.test(source) ||
     /import\s+["']\.\/(?:core\/)?ai["']/.test(source)
   );
 }
 
+/**
+ * True when prior AI setup left unusable stubs (e.g. `local` + `summarizeNote`
+ * without `smart` / without an `ai` import — the old comment-as-configured bug).
+ *
+ * @param source - Existing TypeScript
+ */
+export function isIncompleteAiSetup(source: string): boolean {
+  const hasAiImport = /import\s*\{[^}]*\bai\b[^}]*\}\s*from\s*["']okengine["']/.test(source);
+  const hasSmart = /\bexport\s+const\s+smart\s*=\s*ai\.model\s*\(/.test(source);
+  const hasLocal = /\bexport\s+const\s+local\s*=\s*ai\.model\s*\(/.test(source);
+  const hasSummarize = /\bexport\s+const\s+summarizeNote\s*=/.test(source);
+  const hasAnyModel =
+    /\bexport\s+const\s+(?:smart|local|vision|embedModel)\s*=\s*ai\.model\s*\(/.test(source);
+  if (hasAnyModel && !hasAiImport) return true;
+  if ((hasLocal || hasSummarize) && !hasSmart) return true;
+  return false;
+}
+
+/**
+ * Remove partial AI exports so {@link mergeAiIntoCore} can rewrite a full set.
+ *
+ * @param source - Existing TypeScript
+ */
+export function stripIncompleteAiExports(source: string): string {
+  let next = source;
+  next = next.replace(
+    /(?:\/\*\*[^*]*\*+(?:[^/*][^*]*\*+)*\/\s*)?export\s+const\s+local\s*=\s*ai\.model\s*\(\s*["']local["']\s*,\s*\{[\s\S]*?\}\s*\)\s*;\s*/g,
+    "",
+  );
+  next = next.replace(
+    /(?:\/\*\*[^*]*\*+(?:[^/*][^*]*\*+)*\/\s*)?export\s+const\s+summarizeNote\s*=\s*smart\.prompt\s*\([\s\S]*?\}\s*\)\s*;\s*/g,
+    "",
+  );
+  if (!/\bexport\s+const\s+smart\s*=\s*ai\.model\s*\(/.test(next)) {
+    next = next.replace(
+      /(?:\/\*\*[^*]*\*+(?:[^/*][^*]*\*+)*\/\s*)?export\s+const\s+(?:vision|embedModel)\s*=\s*ai\.model\s*\([\s\S]*?\}\s*\)\s*;\s*/g,
+      "",
+    );
+    next = next.replace(
+      /(?:\/\*\*[^*]*\*+(?:[^/*][^*]*\*+)*\/\s*)?export\s+const\s+docsEmbed\s*=\s*ai\.embed\s*\([\s\S]*?\}\s*\)\s*;\s*/g,
+      "",
+    );
+  }
+  return next.replace(/\n{3,}/g, "\n\n");
+}
+
 /**
  * Append the advanced Notes `summarize-note` prompt when a `smart` model
  * exists but the prompt was never declared (common after older `--ai` runs).
@@ -319,24 +559,26 @@ function hasAiModels(source: string): boolean {
  */
 export function ensureSummarizeNotePrompt(source: string): string {
   let next = source;
-  if (
-    !/\bai\.model\s*\(\s*["']local["']/.test(next) &&
-    /\bai\.model\s*\(\s*["']smart["']/.test(next)
-  ) {
+  const hasSmartExport = /\bexport\s+const\s+smart\s*=/.test(next);
+  const hasLocalExport = /\bexport\s+const\s+local\s*=/.test(next);
+  if (!hasLocalExport && hasSmartExport) {
     next = `${next.trimEnd()}
 
-/** Local inference binding (docker llama.cpp / Ollama via \`OKE_AI_URL\`). */
+/** Local OpenAI-compatible binding (via \`OKE_AI_URL\`). */
 export const local = ai.model("local", {
   provider: "openai-compatible",
-  model: process.env.OKE_AI_LOCAL_MODEL ?? "granite3.3:2b",
+  model: process.env.OKE_AI_LOCAL_MODEL ?? "local-model",
   ...(process.env.OKE_AI_URL?.trim() ? { baseUrl: process.env.OKE_AI_URL.trim() } : {}),
 });
 `;
   }
-  if (/summarize-note/.test(next) || /summarizeNote/.test(next)) {
-    return next;
+  const hasSummarizeExport =
+    /\bexport\s+const\s+summarizeNote\s*=/.test(next) ||
+    /\.prompt\s*\(\s*["']summarize-note["']/.test(next);
+  if (hasSummarizeExport) {
+    return next === source ? next : ensureNamedOkengineImport(next, "ai");
   }
-  if (!/\bai\.model\s*\(\s*["']smart["']/.test(next)) {
+  if (!hasSmartExport) {
     return next;
   }
   const prompt = `
@@ -346,7 +588,7 @@ export const summarizeNote = smart.prompt("summarize-note", {
   timeout: "30s",
 });
 `;
-  return `${next.trimEnd()}\n${prompt}\n`;
+  return ensureNamedOkengineImport(`${next.trimEnd()}\n${prompt}\n`, "ai");
 }
 
 /**