stack/apps/api/src/prisma/llm-encryption.middleware.spec.ts

/**
 * Tests for LLM Encryption Middleware
 *
 * Tests transparent encryption/decryption of LlmProviderInstance.config.apiKey
 * using OpenBao Transit encryption (TransitKey.LLM_CONFIG).
 */

import { describe, it, expect, beforeAll, beforeEach, vi } from "vitest";
import { VaultService } from "../vault/vault.service";
import { TransitKey } from "../vault/vault.constants";
import { registerLlmEncryptionMiddleware } from "./llm-encryption.middleware";

describe("LlmEncryptionMiddleware", () => {
  let mockPrisma: any;
  let mockVaultService: Partial<VaultService>;
  let middlewareFunction: any;

  beforeAll(() => {
    // Mock VaultService
    mockVaultService = {
      encrypt: vi.fn(async (plaintext: string, _key: TransitKey) => {
        return `vault:v1:${plaintext}`;
      }),
      decrypt: vi.fn(async (ciphertext: string, _key: TransitKey) => {
        if (ciphertext.startsWith("vault:v1:")) {
          return ciphertext.replace("vault:v1:", "");
        }
        throw new Error("Invalid ciphertext format");
      }),
    };

    // Create a mock Prisma client
    mockPrisma = {
      $use: vi.fn((fn) => {
        middlewareFunction = fn;
      }),
    };

    // Register the middleware
    registerLlmEncryptionMiddleware(mockPrisma, mockVaultService as VaultService);
  });

  beforeEach(() => {
    // Clear mock call history before each test
    vi.clearAllMocks();
  });

  /**
   * Helper function to call middleware with mock params
   */
  async function callMiddleware(params: any) {
    if (!middlewareFunction) {
      throw new Error("Middleware not registered");
    }

    // Call middleware with a mock next function
    // For write operations, returns params (to check encryption)
    // For read operations, returns mock result data
    return middlewareFunction(params, async (p: any) => {
      if (p.action === "create") {
        // Simulate database returning created record with encrypted data
        return { id: "test-id", ...p.args.data };
      } else if (p.action === "update") {
        return { id: "test-id", ...p.args.data };
      } else if (p.action === "findUnique" || p.action === "findFirst") {
        // Return the mock result for decryption
        return p.mockResult;
      } else if (p.action === "findMany") {
        // Return the mock results array for decryption
        return p.mockResults || [];
      } else if (p.action === "upsert") {
        // Simulate upsert creating new record
        return { id: "test-id", ...p.args.create };
      }
      return p;
    });
  }

  describe("Encryption on create", () => {
    it("should encrypt apiKey when creating new LlmProviderInstance", async () => {
      // Given: New provider config with plaintext apiKey
      const config = {
        endpoint: "https://api.openai.com/v1",
        apiKey: "sk-test-12345",
        organization: "org-test",
      };

      const mockParams = {
        model: "LlmProviderInstance",
        action: "create" as const,
        args: {
          data: {
            providerType: "openai",
            displayName: "Test OpenAI",
            config,
          },
        },
      };

      // When: Middleware processes create
      const result = await callMiddleware(mockParams);

      // Then: VaultService.encrypt called with apiKey and LLM_CONFIG key
      expect(mockVaultService.encrypt).toHaveBeenCalledWith("sk-test-12345", TransitKey.LLM_CONFIG);

      // Then: Config has encrypted apiKey
      expect(result.config.apiKey).toBe("vault:v1:sk-test-12345");
      expect(result.config.endpoint).toBe("https://api.openai.com/v1");
      expect(result.config.organization).toBe("org-test");
    });

    it("should preserve other config fields while encrypting apiKey", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "create" as const,
        args: {
          data: {
            providerType: "claude",
            displayName: "Test Claude",
            config: {
              endpoint: "https://api.anthropic.com",
              apiKey: "sk-ant-secret",
              timeout: 30000,
              maxTokens: 4096,
            },
          },
        },
      };

      const result = await callMiddleware(mockParams);

      expect(result.config.endpoint).toBe("https://api.anthropic.com");
      expect(result.config.apiKey).toBe("vault:v1:sk-ant-secret");
      expect(result.config.timeout).toBe(30000);
      expect(result.config.maxTokens).toBe(4096);
    });

    it("should handle null apiKey gracefully", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "create" as const,
        args: {
          data: {
            providerType: "ollama",
            displayName: "Test Ollama",
            config: {
              endpoint: "http://localhost:11434",
              model: "llama3",
            },
          },
        },
      };

      await callMiddleware(mockParams);

      // Then: No encryption attempted
      expect(mockVaultService.encrypt).not.toHaveBeenCalled();
    });

    it("should handle config with missing apiKey field", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "create" as const,
        args: {
          data: {
            providerType: "ollama",
            displayName: "Test Ollama",
            config: {
              endpoint: "http://localhost:11434",
            },
          },
        },
      };

      const result = await callMiddleware(mockParams);

      expect(mockVaultService.encrypt).not.toHaveBeenCalled();
      expect(result.config.endpoint).toBe("http://localhost:11434");
    });
  });

  describe("Decryption on read", () => {
    it("should decrypt apiKey when reading LlmProviderInstance", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "findUnique" as const,
        args: {},
        mockResult: {
          id: "test-id",
          providerType: "openai",
          displayName: "Test Provider",
          config: {
            endpoint: "https://api.openai.com/v1",
            apiKey: "vault:v1:sk-original-key",
          },
        },
      };

      const result = await callMiddleware(mockParams);

      expect(mockVaultService.decrypt).toHaveBeenCalledWith(
        "vault:v1:sk-original-key",
        TransitKey.LLM_CONFIG
      );
      expect(result.config.apiKey).toBe("sk-original-key");
      expect(result.config.endpoint).toBe("https://api.openai.com/v1");
    });

    it("should decrypt apiKey for all instances in findMany", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "findMany" as const,
        args: {},
        mockResults: [
          {
            id: "id-1",
            providerType: "openai",
            displayName: "OpenAI",
            config: { apiKey: "vault:v1:sk-key-1", endpoint: "https://api.openai.com/v1" },
          },
          {
            id: "id-2",
            providerType: "claude",
            displayName: "Claude",
            config: { apiKey: "vault:v1:sk-ant-key-2", endpoint: "https://api.anthropic.com" },
          },
        ],
      };

      const results = await callMiddleware(mockParams);

      expect(mockVaultService.decrypt).toHaveBeenCalledTimes(2);
      expect(mockVaultService.decrypt).toHaveBeenCalledWith(
        "vault:v1:sk-key-1",
        TransitKey.LLM_CONFIG
      );
      expect(mockVaultService.decrypt).toHaveBeenCalledWith(
        "vault:v1:sk-ant-key-2",
        TransitKey.LLM_CONFIG
      );

      expect(results[0].config.apiKey).toBe("sk-key-1");
      expect(results[1].config.apiKey).toBe("sk-ant-key-2");
    });

    it("should handle plaintext apiKey for backward compatibility", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "findUnique" as const,
        args: {},
        mockResult: {
          id: "test-id",
          providerType: "openai",
          displayName: "Legacy Provider",
          config: {
            endpoint: "https://api.openai.com/v1",
            apiKey: "sk-plaintext-key", // No vault: prefix
          },
        },
      };

      const result = await callMiddleware(mockParams);

      // Then: No decryption attempted (plaintext detected)
      expect(mockVaultService.decrypt).not.toHaveBeenCalled();

      // Then: Plaintext apiKey returned as-is
      expect(result.config.apiKey).toBe("sk-plaintext-key");
    });

    it("should handle missing apiKey gracefully on read", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "findUnique" as const,
        args: {},
        mockResult: {
          id: "test-id",
          providerType: "ollama",
          displayName: "Ollama",
          config: {
            endpoint: "http://localhost:11434",
            // No apiKey
          },
        },
      };

      const result = await callMiddleware(mockParams);

      expect(mockVaultService.decrypt).not.toHaveBeenCalled();
      expect(result.config.endpoint).toBe("http://localhost:11434");
    });
  });

  describe("Idempotent encryption", () => {
    it("should not double-encrypt already encrypted apiKey on update", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "update" as const,
        args: {
          data: {
            config: {
              endpoint: "https://api.openai.com/v1",
              apiKey: "vault:v1:sk-original-key", // Already encrypted
            },
          },
        },
      };

      await callMiddleware(mockParams);

      // Then: No encryption (already encrypted)
      expect(mockVaultService.encrypt).not.toHaveBeenCalled();
    });

    it("should encrypt new plaintext apiKey on update", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "update" as const,
        args: {
          data: {
            config: {
              endpoint: "https://api.openai.com/v1",
              apiKey: "sk-new-key", // Plaintext
            },
          },
        },
      };

      const result = await callMiddleware(mockParams);

      expect(mockVaultService.encrypt).toHaveBeenCalledWith("sk-new-key", TransitKey.LLM_CONFIG);
      expect(result.config.apiKey).toBe("vault:v1:sk-new-key");
    });
  });

  describe("Error handling", () => {
    it("should throw user-facing error when decryption fails", async () => {
      // Mock decryption failure
      vi.spyOn(mockVaultService, "decrypt").mockRejectedValueOnce(new Error("OpenBao unavailable"));

      const mockParams = {
        model: "LlmProviderInstance",
        action: "findUnique" as const,
        args: {},
        mockResult: {
          id: "test-id",
          providerType: "openai",
          displayName: "Test Provider",
          config: {
            endpoint: "https://api.openai.com/v1",
            apiKey: "vault:v1:sk-test-key",
          },
        },
      };

      await expect(callMiddleware(mockParams)).rejects.toThrow(
        /Failed to decrypt LLM provider configuration/
      );
    });
  });

  describe("Upsert operations", () => {
    it("should encrypt apiKey on upsert create", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "upsert" as const,
        args: {
          create: {
            providerType: "openai",
            displayName: "Upserted Provider",
            config: {
              endpoint: "https://api.openai.com/v1",
              apiKey: "sk-upsert-create-key",
            },
          },
          update: {},
        },
      };

      await callMiddleware(mockParams);

      expect(mockVaultService.encrypt).toHaveBeenCalledWith(
        "sk-upsert-create-key",
        TransitKey.LLM_CONFIG
      );
    });

    it("should encrypt apiKey on upsert update", async () => {
      const mockParams = {
        model: "LlmProviderInstance",
        action: "upsert" as const,
        args: {
          create: {
            providerType: "openai",
            displayName: "Should Not Use",
            config: { apiKey: "sk-should-not-use" },
          },
          update: {
            config: {
              endpoint: "https://api.openai.com/v1",
              apiKey: "sk-upsert-update-key",
            },
          },
        },
      };

      await callMiddleware(mockParams);

      // Both create and update paths are encrypted
      expect(mockVaultService.encrypt).toHaveBeenCalledWith(
        "sk-should-not-use",
        TransitKey.LLM_CONFIG
      );
      expect(mockVaultService.encrypt).toHaveBeenCalledWith(
        "sk-upsert-update-key",
        TransitKey.LLM_CONFIG
      );
    });
  });

  describe("Non-LlmProviderInstance models", () => {
    it("should skip encryption for other models", async () => {
      const mockParams = {
        model: "User",
        action: "create" as const,
        args: {
          data: {
            email: "[email protected]",
            name: "Test User",
          },
        },
      };

      await callMiddleware(mockParams);

      expect(mockVaultService.encrypt).not.toHaveBeenCalled();
    });
  });
});