stack/apps/api/src/prisma/llm-encryption.extension.ts

/**
 * LLM Encryption Extension
 *
 * Prisma Client Extension that transparently encrypts/decrypts LLM provider API keys
 * in the LlmProviderInstance.config JSON field using VaultService.
 *
 * Migrated from $use() middleware (removed in Prisma 5+) to $extends() API.
 */

import { Logger } from "@nestjs/common";
import { Prisma } from "@prisma/client";
import type { VaultService } from "../vault/vault.service";
import { TransitKey } from "../vault/vault.constants";

interface LlmProviderInstanceData extends Record<string, unknown> {
  config?: LlmProviderConfig;
}

interface LlmProviderConfig {
  apiKey?: string | null;
  endpoint?: string;
  [key: string]: unknown;
}

export function createLlmEncryptionExtension(vaultService: VaultService) {
  const logger = new Logger("LlmEncryptionExtension");

  return Prisma.defineExtension({
    name: "llm-encryption",
    query: {
      llmProviderInstance: {
        async create({ args, query }) {
          await encryptConfig(args.data as LlmProviderInstanceData, vaultService);
          return query(args);
        },
        async update({ args, query }) {
          await encryptConfig(args.data as LlmProviderInstanceData, vaultService);
          return query(args);
        },
        async updateMany({ args, query }) {
          await encryptConfig(args.data as LlmProviderInstanceData, vaultService);
          return query(args);
        },
        async upsert({ args, query }) {
          await encryptConfig(args.create as LlmProviderInstanceData, vaultService);
          await encryptConfig(args.update as LlmProviderInstanceData, vaultService);
          return query(args);
        },
        async findUnique({ args, query }) {
          const result = await query(args);
          if (result && typeof result === "object") {
            await decryptConfig(result as LlmProviderInstanceData, vaultService, logger);
          }
          return result;
        },
        async findFirst({ args, query }) {
          const result = await query(args);
          if (result && typeof result === "object") {
            await decryptConfig(result as LlmProviderInstanceData, vaultService, logger);
          }
          return result;
        },
        async findMany({ args, query }) {
          const results = await query(args);
          for (const instance of results) {
            await decryptConfig(instance as LlmProviderInstanceData, vaultService, logger);
          }
          return results;
        },
      },
    },
  });
}

async function encryptConfig(
  data: LlmProviderInstanceData,
  vaultService: VaultService
): Promise<void> {
  if (!data.config || typeof data.config !== "object") return;
  const config = data.config;
  if (!config.apiKey || typeof config.apiKey !== "string") return;
  if (isEncrypted(config.apiKey)) return;

  config.apiKey = await vaultService.encrypt(config.apiKey, TransitKey.LLM_CONFIG);
}

async function decryptConfig(
  instance: LlmProviderInstanceData,
  vaultService: VaultService,
  _logger: Logger
): Promise<void> {
  if (!instance.config || typeof instance.config !== "object") return;
  const config = instance.config;
  if (!config.apiKey || typeof config.apiKey !== "string") return;
  if (!isEncrypted(config.apiKey)) return;

  try {
    config.apiKey = await vaultService.decrypt(config.apiKey, TransitKey.LLM_CONFIG);
  } catch (error) {
    const errorMsg = error instanceof Error ? error.message : "Unknown error";
    throw new Error(
      `Failed to decrypt LLM provider configuration. Please re-enter the API key. Details: ${errorMsg}`
    );
  }
}

function isEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") return false;
  if (value.startsWith("vault:v1:")) return true;
  if (isAESEncrypted(value)) return true;
  return false;
}

function isAESEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") return false;
  const parts = value.split(":");
  if (parts.length !== 3) return false;
  return parts.every((part) => /^[0-9a-f]+$/i.test(part));
}