:3

2026-09-21 22:12:37 +02:00
parent b4cb62b490
commit 19b791493f
7 changed files with 425 additions and 10 deletions
@@ -79,6 +79,14 @@ import {
  type ReportReason,
  type ReportStatus,
 } from "./reports.ts";
+import {
+  checkModelAccess,
+  getModelPolicy,
+  initModelUsageTables,
+  policyFor,
+  recordModelUse,
+  setModelPolicyEntry,
+} from "./modelPolicy.ts";
 import type { ConversationRow } from "./db.ts";
 import type { OllamaDriver } from "./drivers/ollama.ts";
 import { DriverRegistry } from "./drivers/registry.ts";
@@ -109,6 +117,7 @@ const ollamaDriver = registry.get("ollama") as OllamaDriver;
 await tools.init();
 initAccessTables();
 initReportsTables();
+initModelUsageTables();
 initTotpTables();

 function json(data: unknown, status = 200): Response {
@@ -980,6 +989,43 @@ const server = Bun.serve<{ ticketId: string | null; isAdmin: boolean }>({
        }
      }

+      // --- admin: per-model access policy (kill-switch + rate limits) ---
+      if (path === "/api/admin/model-policy" && req.method === "GET") {
+        if (!admin) return json({ error: "forbidden" }, 403);
+        if (/^(1|true|yes)$/i.test(url.searchParams.get("refresh") ?? "")) {
+          registry.invalidateModelsCache();
+        }
+        // Full list (policy never hides models from admins) merged with policy.
+        const models = await registry.listAllModelsCached();
+        const policy = getModelPolicy();
+        return json({
+          models: models.map((m) => ({ ...m, ...policyFor(policy, m.id) })),
+        });
+      }
+
+      if (path === "/api/admin/model-policy" && req.method === "PATCH") {
+        if (!admin) return json({ error: "forbidden" }, 403);
+        const { ok, body } = await readJson(req);
+        if (!ok) return json({ error: "invalid JSON" }, 400);
+        const b = body as { model?: unknown; enabled?: unknown; hourly?: unknown; daily?: unknown };
+        if (typeof b.model !== "string" || b.model.length < 3 || b.model.length > 200 || !b.model.includes(":")) {
+          return json({ error: 'model must look like "driver:model"' }, 400);
+        }
+        if (b.enabled === undefined && b.hourly === undefined && b.daily === undefined) {
+          return json({ error: "nothing to update (enabled, hourly, daily)" }, 400);
+        }
+        try {
+          const entry = setModelPolicyEntry(b.model, {
+            enabled: b.enabled as boolean | undefined,
+            hourly: b.hourly as number | undefined,
+            daily: b.daily as number | undefined,
+          });
+          return json({ model: { id: b.model, ...entry } });
+        } catch {
+          return json({ error: "enabled must be boolean, hourly/daily integers 0-1000000" }, 400);
+        }
+      }
+
      // --- admin: access-request wishlist triage ---
      if (path === "/api/admin/access" && req.method === "GET") {
        if (!admin) return json({ error: "forbidden" }, 403);
@@ -1186,7 +1232,11 @@ const server = Bun.serve<{ ticketId: string | null; isAdmin: boolean }>({
          url.searchParams.get("refresh") ?? url.searchParams.get("force") ?? "",
        );
        const models = await registry.listAllModelsCached(wantRefresh && admin);
-        return json({ models });
+        if (admin) return json({ models });
+        // Users never see admin-disabled models (nor select them — /api/chat
+        // enforces the same policy server-side).
+        const policy = getModelPolicy();
+        return json({ models: models.filter((m) => policyFor(policy, m.id).enabled) });
      }

      // --- platform tools (uploads always go through tools) ---
@@ -1237,6 +1287,13 @@ const server = Bun.serve<{ ticketId: string | null; isAdmin: boolean }>({
        } catch (e) {
          return json({ error: (e as Error).message }, 400);
        }
+        // Per-model access policy (kill-switch + per-user rate limits).
+        // Admins bypass; attempts count even when generation later fails.
+        if (!user.isAdmin) {
+          const access = checkModelAccess(user.id, body.model);
+          if (!access.ok) return json({ error: access.message }, access.code);
+          recordModelUse(user.id, body.model);
+        }
        const messages = body.messages as ChatMessage[];
        // Optional persistence into a server-side conversation (must belong to the user).
        let conv = null;
@@ -0,0 +1,132 @@
+import { getDb, getSetting, setSetting } from "./db.ts";
+
+/**
+ * Per-model access policy (admin kill-switch + per-user rate limits).
+ * Stored as one JSON blob in the `settings` table (`model_policy` key);
+ * usage counters live in `model_usage` (hourly + daily fixed windows).
+ *
+ * Policy is global, not per-user; admins always bypass checks.
+ * Limits of 0 mean "unlimited". Models without an entry use the default
+ * (enabled, no limits); entries equal to the default are pruned.
+ */
+
+export interface ModelPolicyEntry {
+  enabled: boolean;
+  /** Max requests per UTC hour (0 = unlimited). */
+  hourly: number;
+  /** Max requests per UTC day (0 = unlimited). */
+  daily: number;
+}
+
+export const DEFAULT_POLICY: ModelPolicyEntry = { enabled: true, hourly: 0, daily: 0 };
+export type ModelPolicy = Record<string, ModelPolicyEntry>;
+
+const SETTINGS_KEY = "model_policy";
+export const MAX_LIMIT = 1_000_000;
+
+function numOrZero(v: unknown): number {
+  return typeof v === "number" && Number.isInteger(v) && v >= 0 && v <= MAX_LIMIT ? v : 0;
+}
+
+function sanitizeEntry(v: unknown): ModelPolicyEntry {
+  const o = (v ?? {}) as Record<string, unknown>;
+  return {
+    enabled: typeof o.enabled === "boolean" ? o.enabled : true,
+    hourly: numOrZero(o.hourly),
+    daily: numOrZero(o.daily),
+  };
+}
+
+function isDefault(e: ModelPolicyEntry): boolean {
+  return e.enabled && e.hourly === 0 && e.daily === 0;
+}
+
+export function getModelPolicy(): ModelPolicy {
+  const raw = getSetting(SETTINGS_KEY, "");
+  if (!raw) return {};
+  try {
+    const obj = JSON.parse(raw) as Record<string, unknown>;
+    if (!obj || typeof obj !== "object") return {};
+    const out: ModelPolicy = {};
+    for (const [k, v] of Object.entries(obj)) {
+      if (k && k.length <= 200) out[k] = sanitizeEntry(v);
+    }
+    return out;
+  } catch {
+    return {};
+  }
+}
+
+export function policyFor(policy: ModelPolicy, modelId: string): ModelPolicyEntry {
+  return policy[modelId] ?? { ...DEFAULT_POLICY };
+}
+
+export function setModelPolicyEntry(
+  modelId: string,
+  patch: { enabled?: boolean; hourly?: number; daily?: number },
+): ModelPolicyEntry {
+  if (patch.enabled !== undefined && typeof patch.enabled !== "boolean") throw new Error("BAD_POLICY");
+  if (patch.hourly !== undefined && !Number.isInteger(patch.hourly)) throw new Error("BAD_POLICY");
+  if (patch.daily !== undefined && !Number.isInteger(patch.daily)) throw new Error("BAD_POLICY");
+  const policy = getModelPolicy();
+  const cur = policy[modelId] ?? { ...DEFAULT_POLICY };
+  const next: ModelPolicyEntry = {
+    enabled: patch.enabled ?? cur.enabled,
+    hourly: patch.hourly ?? cur.hourly,
+    daily: patch.daily ?? cur.daily,
+  };
+  if (next.hourly < 0 || next.hourly > MAX_LIMIT || next.daily < 0 || next.daily > MAX_LIMIT) {
+    throw new Error("BAD_POLICY");
+  }
+  if (isDefault(next)) delete policy[modelId];
+  else policy[modelId] = next;
+  setSetting(SETTINGS_KEY, JSON.stringify(policy));
+  return { ...DEFAULT_POLICY, ...policy[modelId] };
+}
+
+export function initModelUsageTables(): void {
+  getDb().query(`CREATE TABLE IF NOT EXISTS model_usage (
+    user_id TEXT NOT NULL,
+    model TEXT NOT NULL,
+    window TEXT NOT NULL,
+    bucket INTEGER NOT NULL,
+    count INTEGER NOT NULL DEFAULT 0,
+    PRIMARY KEY (user_id, model, window, bucket)
+  )`).run();
+}
+
+function buckets(now = Date.now()): { h: number; d: number } {
+  return { h: Math.floor(now / 3600_000), d: Math.floor(now / 86400_000) };
+}
+
+function windowCount(userId: string, modelId: string, window: "h" | "d", bucket: number): number {
+  const row = getDb().query(
+    "SELECT count AS n FROM model_usage WHERE user_id = ? AND model = ? AND window = ? AND bucket = ?",
+  ).get(userId, modelId, window, bucket) as { n: number } | null;
+  return row?.n ?? 0;
+}
+
+export type ModelAccess = { ok: true } | { ok: false; code: 403 | 429; message: string };
+
+export function checkModelAccess(userId: string, modelId: string): ModelAccess {
+  const p = policyFor(getModelPolicy(), modelId);
+  if (!p.enabled) return { ok: false, code: 403, message: "model disabled by administrator" };
+  const { h, d } = buckets();
+  if (p.hourly > 0 && windowCount(userId, modelId, "h", h) >= p.hourly) {
+    return { ok: false, code: 429, message: "hourly limit reached for this model, try again later" };
+  }
+  if (p.daily > 0 && windowCount(userId, modelId, "d", d) >= p.daily) {
+    return { ok: false, code: 429, message: "daily limit reached for this model, try again tomorrow" };
+  }
+  return { ok: true };
+}
+
+export function recordModelUse(userId: string, modelId: string): void {
+  const d = getDb();
+  const { h, d: day } = buckets();
+  d.query(`INSERT INTO model_usage (user_id, model, window, bucket, count) VALUES (?, ?, 'h', ?, 1)
+    ON CONFLICT(user_id, model, window, bucket) DO UPDATE SET count = count + 1`).run(userId, modelId, h);
+  d.query(`INSERT INTO model_usage (user_id, model, window, bucket, count) VALUES (?, ?, 'd', ?, 1)
+    ON CONFLICT(user_id, model, window, bucket) DO UPDATE SET count = count + 1`).run(userId, modelId, day);
+  d.query("DELETE FROM model_usage WHERE (window = 'h' AND bucket < ?) OR (window = 'd' AND bucket < ?)").run(h - 1, day - 1);
+}