chore: import upstream snapshot with attribution

apps/sim/lib/api-key/byok.test.ts

@@ -0,0 +1,162 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockOrderBy, mockDecryptSecret } = vi.hoisted(() => ({
+  mockOrderBy: vi.fn(),
+  mockDecryptSecret: vi.fn(),
+}))
+
+vi.mock('@sim/db', () => ({
+  db: {
+    select: vi.fn(() => ({
+      from: vi.fn(() => ({
+        where: vi.fn(() => ({ orderBy: mockOrderBy })),
+      })),
+    })),
+  },
+}))
+
+vi.mock('@/lib/core/security/encryption', () => ({
+  decryptSecret: mockDecryptSecret,
+}))
+
+vi.mock('@/lib/core/config/api-keys', () => ({
+  getRotatingApiKey: vi.fn(),
+}))
+
+vi.mock('@/lib/core/config/env', () => ({
+  env: {},
+}))
+
+vi.mock('@/lib/core/config/env-flags', () => ({
+  isHosted: false,
+}))
+
+vi.mock('@/providers/models', () => ({
+  getProviderFileAttachment: vi
+    .fn()
+    .mockReturnValue({ maxBytes: 10 * 1024 * 1024, strategy: 'inline' }),
+  INLINE_ATTACHMENT_MAX_BYTES: 10 * 1024 * 1024,
+  getHostedModels: vi.fn(() => []),
+}))
+
+vi.mock('@/providers/utils', () => ({
+  PROVIDER_PLACEHOLDER_KEY: 'placeholder',
+}))
+
+vi.mock('@/stores/providers/store', () => ({
+  useProvidersStore: { getState: vi.fn() },
+}))
+
+import { getBYOKKey } from '@/lib/api-key/byok'
+
+/**
+ * Rotation counters in the module under test are keyed by
+ * `${workspaceId}:${providerId}` and persist for the process lifetime, so
+ * each test uses a unique workspace id to start from a fresh cursor.
+ */
+let testIndex = 0
+const uniqueWorkspaceId = () => `workspace-${++testIndex}`
+
+const storedKey = (id: string) => ({ id, encryptedApiKey: `encrypted-${id}` })
+
+describe('getBYOKKey', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockOrderBy.mockResolvedValue([])
+    mockDecryptSecret.mockImplementation(async (encrypted: string) => ({
+      decrypted: encrypted.replace('encrypted-', 'decrypted-'),
+    }))
+  })
+
+  it('returns null when no workspaceId is provided', async () => {
+    expect(await getBYOKKey(undefined, 'openai')).toBeNull()
+    expect(await getBYOKKey(null, 'openai')).toBeNull()
+  })
+
+  it('returns null when the workspace has no keys for the provider', async () => {
+    expect(await getBYOKKey(uniqueWorkspaceId(), 'openai')).toBeNull()
+  })
+
+  it('returns the same key on every call when only one key is stored', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1')])
+
+    for (let call = 0; call < 3; call++) {
+      expect(await getBYOKKey(workspaceId, 'openai')).toEqual({
+        apiKey: 'decrypted-key-1',
+        isBYOK: true,
+      })
+    }
+  })
+
+  it('round-robins across multiple keys in creation order', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1'), storedKey('key-2'), storedKey('key-3')])
+
+    const apiKeys = []
+    for (let call = 0; call < 4; call++) {
+      const result = await getBYOKKey(workspaceId, 'openai')
+      apiKeys.push(result?.apiKey)
+    }
+
+    expect(apiKeys).toEqual([
+      'decrypted-key-1',
+      'decrypted-key-2',
+      'decrypted-key-3',
+      'decrypted-key-1',
+    ])
+  })
+
+  it('reads the key list fresh from the database on every call', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1')])
+
+    await getBYOKKey(workspaceId, 'openai')
+    await getBYOKKey(workspaceId, 'openai')
+    await getBYOKKey(workspaceId, 'openai')
+
+    expect(mockOrderBy).toHaveBeenCalledTimes(3)
+  })
+
+  it('tracks rotation independently per provider within a workspace', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1'), storedKey('key-2')])
+
+    expect((await getBYOKKey(workspaceId, 'openai'))?.apiKey).toBe('decrypted-key-1')
+    expect((await getBYOKKey(workspaceId, 'anthropic'))?.apiKey).toBe('decrypted-key-1')
+    expect((await getBYOKKey(workspaceId, 'openai'))?.apiKey).toBe('decrypted-key-2')
+  })
+
+  it('skips a key that fails to decrypt and returns the next one', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1'), storedKey('key-2')])
+    mockDecryptSecret.mockImplementation(async (encrypted: string) => {
+      if (encrypted === 'encrypted-key-1') {
+        throw new Error('corrupt ciphertext')
+      }
+      return { decrypted: encrypted.replace('encrypted-', 'decrypted-') }
+    })
+
+    expect(await getBYOKKey(workspaceId, 'openai')).toEqual({
+      apiKey: 'decrypted-key-2',
+      isBYOK: true,
+    })
+  })
+
+  it('returns null when every key fails to decrypt', async () => {
+    const workspaceId = uniqueWorkspaceId()
+    mockOrderBy.mockResolvedValue([storedKey('key-1'), storedKey('key-2')])
+    mockDecryptSecret.mockRejectedValue(new Error('corrupt ciphertext'))
+
+    expect(await getBYOKKey(workspaceId, 'openai')).toBeNull()
+  })
+
+  it('returns null when the keys query throws', async () => {
+    mockOrderBy.mockRejectedValue(new Error('database unavailable'))
+
+    expect(await getBYOKKey(uniqueWorkspaceId(), 'openai')).toBeNull()
+  })
+})