feat(settings): power Loki with your own model — paste a key, get your best model

.env.example

@@ -66,6 +66,13 @@ APP_LOCALE=de-CH
 # Docker only — Postgres password (must match DATABASE_URL above)
 POSTGRES_PASSWORD=changeme
 
+# "Power Loki with your own model" (Settings → AI). Seals each user's API key
+# at rest (@bitbaum/ai-kit/seal, AES-256-GCM). Without it the feature says it
+# is not switched on and chat keeps using the free chain. Generate once with
+# `openssl rand -base64 32`; ROTATING it makes every stored key unreadable
+# (users are then asked to paste their key again).
+BYOK_SEAL_SECRET=
+
 # ── Brand alias scheme ───────────────────────────────────────────────────────
 # Every env var below is read as APP_<KEY> first, then COCKPIT_<KEY> as a
 # legacy fallback (mirror of scripts/_brand.sh:_brand_env). Set whichever

drizzle/0081_user_model_keys.sql

@@ -0,0 +1,26 @@
+-- Migration: user_model_keys — the model a user brings to power Loki.
+--
+-- WHY THIS EXISTS
+-- Loki's own chat ran only on the shared, rationed free chain: when the day's
+-- budget ran out a user was told to wait, and there was no way to say "use my
+-- Anthropic / OpenAI / OpenRouter / … key instead". This stores that choice:
+-- a vendor from ai-kit's closed list, the model picked there, and the key.
+--
+-- THE KEY IS STORED SEALED (ai-kit/seal: AES-256-GCM, scrypt-derived from the
+-- app's BYOK_SEAL_SECRET, `iv:tag:ciphertext`). The database alone cannot read
+-- it, no API ever returns it, and `key_hint` ("…abcd") is the only part shown.
+-- One row per user; removing the user removes the key.
+--
+-- Hand-written, like 0073–0080: drizzle/meta's snapshot stops at 0072, so
+-- `drizzle-kit generate` diffs against a stale picture and asks whether this
+-- table is a rename of beacon_sessions.
+CREATE TABLE IF NOT EXISTS user_model_keys (
+  user_id uuid PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE,
+  vendor text NOT NULL,
+  model text NOT NULL,
+  sealed_key text NOT NULL,
+  key_hint text NOT NULL,
+  verified_at timestamptz NOT NULL,
+  created_at timestamptz NOT NULL DEFAULT now(),
+  updated_at timestamptz NOT NULL DEFAULT now()
+);

drizzle/meta/_journal.json

@@ -568,6 +568,13 @@
       "when": 1767225680000,
       "tag": "0080_project_invitations",
       "breakpoints": true
+    },
+    {
+      "idx": 81,
+      "version": "7",
+      "when": 1767225681000,
+      "tag": "0081_user_model_keys",
+      "breakpoints": true
     }
   ]
 }

package.json

@@ -114,7 +114,7 @@
   "dependencies": {
     "@auth/core": "0.41.3",
     "@auth/drizzle-adapter": "^1.11.3",
-    "@bitbaum/ai-kit": "^1.11.0",
+    "@bitbaum/ai-kit": "^1.17.0",
     "@bitbaum/chatkit": "github:bitbaum/chatkit#v0.1.0",
     "@bitbaum/mail-kit": "^0.1.1",
     "@dnd-kit/core": "^6.3.1",

scripts/test/own-model.ts

@@ -0,0 +1,142 @@
+/**
+ * A user's own model: their key goes to their vendor, and only there.
+ *
+ * THE PROPERTIES THIS PINS
+ * ------------------------
+ * Loki's chat can run on a key the user brought (Settings → AI). The walker
+ * that makes the call looks a link's key up BY NAME, so the whole safety of
+ * the feature rests on one rule: a user's link reads its key only from the
+ * per-call env ai-kit's byokChain returns, and a server link reads only
+ * `process.env`. Break it either way and one of two things happens:
+ *
+ *   • a deployment that happens to set BYOK_API_KEY answers every user's turn
+ *     on the operator's account, or
+ *   • a user's key rides along on a server link to a vendor they never chose.
+ *
+ * The test drives the REAL walker (`callModelWithTools`) with a fake fetch and
+ * reads the request it actually sends.
+ *
+ * Run: npx tsx scripts/test/own-model.ts
+ */
+import assert from "node:assert/strict";
+import { byokChain } from "@bitbaum/ai-kit/byok";
+import { callModelWithTools } from "@/lib/agent/llm";
+import { OWN_MODEL_KEY_ENV, isOwnModelLink, keyForLink, ownModelFrom } from "@/lib/own-model";
+import type { ChatLink } from "@/config/chat-models";
+
+let passed = 0;
+async function check(name: string, fn: () => void | Promise<void>) {
+  await fn();
+  passed++;
+  console.log(`  ✓ ${name}`);
+}
+
+const USER_KEY = "sk-ant-test-not-a-real-key-1234";
+const config = { vendor: "anthropic" as const, apiKey: USER_KEY, model: "claude-opus-5.5" };
+
+async function main() {
+  await check("the key name here is the one ai-kit's byokChain actually uses", () => {
+    const { chain, env } = byokChain(config);
+    assert.equal(chain[0]!.provider.keyEnv, OWN_MODEL_KEY_ENV);
+    assert.equal((env as Record<string, string>)[OWN_MODEL_KEY_ENV], USER_KEY);
+  });
+
+  await check("a user's link never falls back to a key the server happens to hold", () => {
+    const own = ownModelFrom(config);
+    const link = own.chain[0]!;
+    process.env[OWN_MODEL_KEY_ENV] = "server-set-this-by-accident";
+    try {
+      assert.equal(keyForLink(link, own), USER_KEY);
+      assert.equal(
+        keyForLink(link, undefined),
+        undefined,
+        "no per-call env → no key, never process.env",
+      );
+    } finally {
+      delete process.env[OWN_MODEL_KEY_ENV];
+    }
+  });
+
+  await check("a server link never takes a key a user supplied", () => {
+    const serverLink = {
+      provider: {
+        id: "groq",
+        baseUrl: "https://api.groq.com/openai/v1",
+        keyEnv: "GROQ_API_KEY",
+        models: ["m"],
+      },
+      model: "m",
+    } as unknown as ChatLink;
+    const hostile = { env: { GROQ_API_KEY: "user-injected" } };
+    process.env.GROQ_API_KEY = "the-server-groq-key";
+    try {
+      assert.equal(isOwnModelLink(serverLink), false);
+      assert.equal(keyForLink(serverLink, hostile), "the-server-groq-key");
+    } finally {
+      delete process.env.GROQ_API_KEY;
+    }
+  });
+
+  await check(
+    "a chat turn on the user's model sends ONE request, to their vendor, with their key and model",
+    async () => {
+      const own = ownModelFrom(config);
+      const sent: { url: string; auth: string | undefined; model: unknown; allHeaders: string }[] =
+        [];
+      const realFetch = globalThis.fetch;
+      process.env.GROQ_API_KEY = "the-server-groq-key";
+      globalThis.fetch = (async (url: string | URL | Request, init?: RequestInit) => {
+        const headers = (init?.headers ?? {}) as Record<string, string>;
+        sent.push({
+          url: String(url),
+          auth: headers.Authorization,
+          model: JSON.parse(String(init?.body ?? "{}")).model,
+          allHeaders: JSON.stringify(headers),
+        });
+        return new Response(
+          JSON.stringify({
+            choices: [{ message: { role: "assistant", content: "Hello from your model." } }],
+            usage: { prompt_tokens: 3, completion_tokens: 5, total_tokens: 8 },
+          }),
+          { status: 200, headers: { "content-type": "application/json" } },
+        );
+      }) as typeof fetch;
+      try {
+        const turn = await callModelWithTools({
+          messages: [{ role: "user", content: "hi" }],
+          tools: [],
+          validToolNames: [],
+          feature: "own-model-test",
+          own,
+        });
+        assert.equal(sent.length, 1, `expected one request, got ${sent.length}`);
+        assert.equal(sent[0]!.url, "https://api.anthropic.com/v1/chat/completions");
+        assert.equal(sent[0]!.auth, `Bearer ${USER_KEY}`);
+        assert.equal(sent[0]!.model, "claude-opus-5.5");
+        assert.ok(!sent[0]!.allHeaders.includes("the-server-groq-key"), "a server key rode along");
+        assert.match(JSON.stringify(turn), /Hello from your model/);
+      } finally {
+        globalThis.fetch = realFetch;
+        delete process.env.GROQ_API_KEY;
+      }
+    },
+  );
+
+  await check("OpenRouter gets its attribution headers; a direct lab does not", () => {
+    assert.ok(
+      ownModelFrom({ ...config, vendor: "openrouter", model: "anthropic/claude-opus-5.5" })
+        .extraHeaders,
+    );
+    assert.equal(ownModelFrom(config).extraHeaders, undefined);
+  });
+
+  await check("the label names the vendor and model, never the key", () => {
+    const label = ownModelFrom(config).label;
+    assert.match(label, /Anthropic · claude-opus-5\.5/);
+    assert.ok(!label.includes(USER_KEY));
+  });
+
+  console.log(`\nown-model: ${passed} passed`);
+}
+
+void main();

src/app/api/settings/model/probe/route.ts

@@ -0,0 +1,55 @@
+/**
+ * POST { vendor, apiKey? } → "does this key work, and what can it use?"
+ *
+ * Nothing is stored. The settings screen calls this the moment a key is
+ * pasted, so it can say whether the key works and offer the models it can
+ * reach with the best one preselected (ai-kit's probeByokKey). With no apiKey
+ * it probes the key already stored for that vendor — so changing model never
+ * means pasting the key again.
+ *
+ * Server-side because vendors refuse browser-origin calls, and because the
+ * hosts it may contact are ai-kit's closed list, never a URL from the request.
+ * Rate-limited per user: it makes one or two outbound calls per request.
+ */
+import type { NextRequest } from "next/server";
+import { NextResponse } from "next/server";
+import { BYOK_VENDOR_IDS, type ByokVendorId } from "@bitbaum/ai-kit/byok";
+import { probeByokKey } from "@bitbaum/ai-kit/byok-probe";
+import { getApiUserId } from "@/lib/session";
+import { jsonError, jsonOk, readJsonBody, z } from "@/lib/api/route-helpers";
+import { checkRateLimit } from "@/lib/rate-limit";
+import { getOwnModel } from "@/db/queries/user-model-keys";
+
+export const runtime = "nodejs";
+
+const Body = z.object({
+  vendor: z.enum(BYOK_VENDOR_IDS as [ByokVendorId, ...ByokVendorId[]]),
+  apiKey: z.string().trim().min(8).max(400).optional(),
+});
+
+export async function POST(req: NextRequest) {
+  const userId = await getApiUserId();
+  if (!userId) return jsonError("Unauthorized", 401);
+  if (!checkRateLimit(`own-model:${userId}`, 20, 10 * 60_000)) {
+    return jsonError("Too many attempts — wait a few minutes and try again.", 429);
+  }
+  const body = await readJsonBody(req, Body);
+  if (body instanceof NextResponse) return body;
+
+  let apiKey = body.apiKey;
+  if (!apiKey) {
+    const stored = await getOwnModel(userId);
+    if (!stored || stored.vendor !== body.vendor) {
+      return jsonError("Paste your key for this provider first.", 400);
+    }
+    apiKey = stored.apiKey;
+  }
+
+  const probe = await probeByokKey(body.vendor, apiKey);
+  return jsonOk({
+    works: probe.ok,
+    message: probe.message,
+    models: probe.models,
+    suggested: probe.suggested,
+  });
+}

src/app/api/settings/model/route.ts

@@ -0,0 +1,103 @@
+/**
+ * The model a user brings to power Loki — read, save, change, remove.
+ *
+ *   GET    → { available, current: { vendor, model, keyHint, verifiedAt } | null }
+ *   PUT    { vendor, model, apiKey? } → saves; apiKey omitted = change model, keep key
+ *   DELETE → forgets the key
+ *
+ * The key is never returned by any method; `keyHint` ("…abcd") is all a client
+ * ever sees. Checking a key WITHOUT saving it is ./probe.
+ */
+import type { NextRequest } from "next/server";
+import { NextResponse } from "next/server";
+import { BYOK_VENDOR_IDS, isByokConfig, type ByokVendorId } from "@bitbaum/ai-kit/byok";
+import { probeByokKey } from "@bitbaum/ai-kit/byok-probe";
+import { getApiUserId } from "@/lib/session";
+import { jsonError, jsonOk, readJsonBody, z } from "@/lib/api/route-helpers";
+import { checkRateLimit } from "@/lib/rate-limit";
+import {
+  deleteOwnModel,
+  getOwnModel,
+  getOwnModelSummary,
+  ownModelSealSecret,
+  saveOwnModel,
+  setOwnModelChoice,
+  type OwnModelSummary,
+} from "@/db/queries/user-model-keys";
+
+export const runtime = "nodejs";
+
+const NOT_SET_UP =
+  "Connecting your own model isn't switched on for this server yet — Loki keeps using its free models.";
+
+function view(summary: OwnModelSummary | null) {
+  return summary
+    ? {
+        vendor: summary.vendor,
+        model: summary.model,
+        keyHint: summary.keyHint,
+        verifiedAt: summary.verifiedAt.toISOString(),
+      }
+    : null;
+}
+
+export async function GET() {
+  const userId = await getApiUserId();
+  if (!userId) return jsonError("Unauthorized", 401);
+  return jsonOk({
+    available: ownModelSealSecret() !== null,
+    current: view(await getOwnModelSummary(userId)),
+  });
+}
+
+const PutBody = z.object({
+  vendor: z.enum(BYOK_VENDOR_IDS as [ByokVendorId, ...ByokVendorId[]]),
+  model: z.string().trim().min(1).max(200),
+  apiKey: z.string().trim().min(8).max(400).optional(),
+});
+
+export async function PUT(req: NextRequest) {
+  const userId = await getApiUserId();
+  if (!userId) return jsonError("Unauthorized", 401);
+  if (!ownModelSealSecret()) return jsonError(NOT_SET_UP, 503);
+  // Saving probes the vendor, so it shares the probe's allowance.
+  if (!checkRateLimit(`own-model:${userId}`, 20, 10 * 60_000)) {
+    return jsonError("Too many attempts — wait a few minutes and try again.", 429);
+  }
+  const body = await readJsonBody(req, PutBody);
+  if (body instanceof NextResponse) return body;
+
+  // Change the model on the key already stored: no key in the request.
+  if (!body.apiKey) {
+    const stored = await getOwnModel(userId);
+    if (!stored || stored.vendor !== body.vendor) {
+      return jsonError("Paste your key for this provider first.", 400);
+    }
+    const probe = await probeByokKey(stored.vendor, stored.apiKey);
+    if (!probe.ok) return jsonError(probe.message, 400, { status: probe.status });
+    if (probe.models.length > 0 && !probe.models.includes(body.model)) {
+      return jsonError(`Your key can't use ${body.model}. Pick one from the list.`, 400);
+    }
+    return jsonOk({ current: view(await setOwnModelChoice(userId, body.model)) });
+  }
+
+  const config = { vendor: body.vendor, model: body.model, apiKey: body.apiKey };
+  if (!isByokConfig(config)) return jsonError("That doesn't look like a key.", 400);
+
+  // Never trust a client's "it worked": the key is checked again, here, before
+  // it is stored — a saved key that does not work is a chat that fails later
+  // with nobody knowing why.
+  const probe = await probeByokKey(config.vendor, config.apiKey);
+  if (!probe.ok) return jsonError(probe.message, 400, { status: probe.status });
+  if (probe.models.length > 0 && !probe.models.includes(config.model)) {
+    return jsonError(`Your key can't use ${config.model}. Pick one from the list.`, 400);
+  }
+  return jsonOk({ current: view(await saveOwnModel(userId, config)) });
+}
+
+export async function DELETE() {
+  const userId = await getApiUserId();
+  if (!userId) return jsonError("Unauthorized", 401);
+  await deleteOwnModel(userId);
+  return jsonOk({ current: null });
+}