ThothII/backend/src/models/runtime-model-catalog.ts

import {
  closeSync, constants, fstatSync, lstatSync, openSync, readFileSync, type Stats,
} from "node:fs";
import { z } from "zod";

const MAX_CATALOG_BYTES = 1024 * 1024;
const RUNTIME_CATALOG_FILE = "/run/thothii-model-catalog/catalog.json";
const canonicalId = z.string().regex(/^[a-z][a-z0-9._-]{0,63}\/[A-Za-z0-9][A-Za-z0-9._:-]{0,255}$/);
const secretBundleKey = /^[A-Z][A-Z0-9_]{0,63}$/;

const endpointSchema = z.object({
  baseUrl: z.string().url(),
  apiVersion: z.string().optional(),
}).strict();

const authenticationSchema = z.object({
  mode: z.enum(["secret_env", "pi_auth", "none"]),
  apiKeyEnv: z.string().optional(),
}).strict();

const runtimeModelSchema = z.object({
  id: canonicalId,
  provider: z.string().min(1),
  model: z.string().min(1),
  label: z.string().min(1),
  upstreamModel: z.string().min(1),
  endpoint: endpointSchema.optional(),
  authentication: authenticationSchema,
  sessionAdapter: z.object({ mode: z.enum(["pi_builtin", "openai_compatible"]) }).strict().optional(),
  metadataAdapter: z.object({ litellmProvider: z.string().min(1) }).strict().optional(),
  session: z.object({
    reasoning: z.boolean(),
    input: z.array(z.string()).optional(),
    cost: z.object({
      input: z.number(), output: z.number(), cacheRead: z.number(), cacheWrite: z.number(),
    }).strict().optional(),
    contextWindow: z.number().int().positive().optional(),
    maxTokens: z.number().int().positive().optional(),
    compatibility: z.object({
      supportsDeveloperRole: z.boolean(),
      supportsReasoningEffort: z.boolean(),
      supportsStore: z.boolean(),
      maxTokensField: z.string().optional(),
    }).strict().optional(),
  }).strict().optional(),
  metadataGeneration: z.object({ disableThinking: z.boolean() }).strict().optional(),
}).strict();

const catalogSchema = z.object({
  schemaVersion: z.literal(1),
  defaultSession: canonicalId,
  defaultMetadataGeneration: canonicalId.optional(),
  embedding: z.object({ id: canonicalId, dimensions: z.number().int().positive() }).strict(),
  models: z.array(runtimeModelSchema).max(64),
}).strict();

export type RuntimeModel = z.infer<typeof runtimeModelSchema>;

export interface RuntimeModelCatalog {
  readonly defaultSession: string | null;
  readonly defaultMetadataGeneration: string | null;
  readonly embedding: Readonly<{ id: string; dimensions: number }> | null;
  sessionModels(): readonly RuntimeModel[];
  metadataModels(): readonly RuntimeModel[];
  hasSession(id: string): boolean;
}

class RestartLoadedRuntimeModelCatalog implements RuntimeModelCatalog {
  readonly defaultSession: string | null;
  readonly defaultMetadataGeneration: string | null;
  readonly embedding: Readonly<{ id: string; dimensions: number }> | null;
  readonly #sessions: readonly RuntimeModel[];
  readonly #metadata: readonly RuntimeModel[];
  readonly #sessionIds: ReadonlySet<string>;

  constructor(catalog?: z.infer<typeof catalogSchema>) {
    this.defaultSession = catalog?.defaultSession ?? null;
    this.defaultMetadataGeneration = catalog?.defaultMetadataGeneration ?? null;
    this.embedding = catalog ? Object.freeze({ ...catalog.embedding }) : null;
    this.#sessions = Object.freeze((catalog?.models ?? []).filter((model) => model.session !== undefined));
    this.#metadata = Object.freeze((catalog?.models ?? []).filter((model) => model.metadataGeneration !== undefined));
    this.#sessionIds = new Set(this.#sessions.map((model) => model.id));
  }

  sessionModels(): readonly RuntimeModel[] { return this.#sessions.map((model) => ({ ...model })); }
  metadataModels(): readonly RuntimeModel[] { return this.#metadata.map((model) => ({ ...model })); }
  hasSession(id: string): boolean { return this.#sessionIds.has(id); }
}

function protectedCatalogStat(file: string, info: Stats): boolean {
  const mode = info.mode & 0o777;
  if (!info.isFile() || info.isSymbolicLink() || info.nlink !== 1
    || info.size < 1 || info.size > MAX_CATALOG_BYTES) return false;
  if (file === RUNTIME_CATALOG_FILE && info.uid === 0 && (mode === 0o444 || mode === 0o644)) return true;
  return info.uid === (process.getuid?.() ?? info.uid) && (mode === 0o400 || mode === 0o600 || mode === 0o644);
}

function readProtectedCatalog(file: string): unknown {
  let descriptor: number | undefined;
  try {
    const before = lstatSync(file);
    if (!protectedCatalogStat(file, before)) throw new Error("runtime model catalog is unavailable");
    descriptor = openSync(file, constants.O_RDONLY | constants.O_NOFOLLOW);
    const opened = fstatSync(descriptor);
    if (!protectedCatalogStat(file, opened)
      || before.dev !== opened.dev || before.ino !== opened.ino) throw new Error("runtime model catalog is unavailable");
    const source = readFileSync(descriptor, "utf8");
    const after = fstatSync(descriptor);
    const current = lstatSync(file);
    if (!protectedCatalogStat(file, after) || !protectedCatalogStat(file, current)
      || opened.dev !== after.dev || opened.ino !== after.ino
      || opened.dev !== current.dev || opened.ino !== current.ino) throw new Error("runtime model catalog is unavailable");
    return JSON.parse(source);
  } catch {
    throw new Error("runtime model catalog is unavailable");
  } finally {
    if (descriptor !== undefined) try { closeSync(descriptor); } catch { /* sanitized above */ }
  }
}

export function loadRuntimeModelCatalog(file?: string): RuntimeModelCatalog {
  if (!file) return new RestartLoadedRuntimeModelCatalog();
  const parsed = catalogSchema.safeParse(readProtectedCatalog(file));
  if (!parsed.success) throw new Error("runtime model catalog is invalid");
  if (parsed.data.models.some((model) => !validRuntimeModel(model))) {
    throw new Error("runtime model catalog is invalid");
  }
  const ids = new Set(parsed.data.models.map((model) => model.id));
  if (ids.size !== parsed.data.models.length) throw new Error("runtime model catalog contains duplicate models");
  const sessions = parsed.data.models.filter((model) => model.session !== undefined).map((model) => model.id);
  const metadata = parsed.data.models.filter((model) => model.metadataGeneration !== undefined).map((model) => model.id);
  if (!sessions.includes(parsed.data.defaultSession)) throw new Error("runtime model catalog session default is invalid");
  if ((metadata.length > 0) !== (parsed.data.defaultMetadataGeneration !== undefined)
    || (parsed.data.defaultMetadataGeneration !== undefined
      && !metadata.includes(parsed.data.defaultMetadataGeneration))) {
    throw new Error("runtime model catalog metadata default is invalid");
  }
  return new RestartLoadedRuntimeModelCatalog(parsed.data);
}

function validRuntimeModel(model: RuntimeModel): boolean {
  if ((model.session !== undefined) !== (model.sessionAdapter !== undefined)) return false;
  if ((model.metadataGeneration !== undefined) !== (model.metadataAdapter !== undefined)) return false;
  switch (model.authentication.mode) {
    case "secret_env":
      return model.authentication.apiKeyEnv !== undefined
        && secretBundleKey.test(model.authentication.apiKeyEnv);
    case "pi_auth":
      return model.authentication.apiKeyEnv === undefined
        && model.metadataGeneration === undefined
        && model.sessionAdapter?.mode === "pi_builtin";
    case "none":
      return model.authentication.apiKeyEnv === undefined && model.endpoint !== undefined;
  }
}

export function splitCanonicalModelId(id: string): { provider: string; model: string } {
  const slash = id.indexOf("/");
  if (slash <= 0 || slash === id.length - 1) throw new Error("model identity is invalid");
  return { provider: id.slice(0, slash), model: id.slice(slash + 1) };
}