simstudioai--sim/apps/sim/app/api/knowledge/utils.test.ts

/**
 * @vitest-environment node
 *
 * Knowledge Utils Unit Tests
 *
 * This file contains unit tests for the knowledge base utility functions,
 * including access checks, document processing, and embedding generation.
 */
import {
  dbChainMockFns,
  defaultMockEnv,
  queueTableRows,
  resetDbChainMock,
  schemaMock,
} from '@sim/testing'
import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
import * as billingAttributionModule from '@/lib/billing/core/billing-attribution'
import { env } from '@/lib/core/config/env'
import * as documentsUtilsModule from '@/lib/knowledge/documents/utils'
import * as workspacesUtilsModule from '@/lib/workspaces/utils'

const envSnapshot = { ...env }

afterAll(() => {
  for (const key of Object.keys(env)) {
    delete (env as Record<string, unknown>)[key]
  }
  Object.assign(env, envSnapshot)
  resetDbChainMock()
  retrySpy.mockRestore()
  vi.mocked(workspacesUtilsModule.getWorkspaceBilledAccountUserId).mockRestore()
  vi.mocked(billingAttributionModule.assertBillingAttributionSnapshot).mockRestore()
  vi.mocked(billingAttributionModule.checkAttributedUsageLimits).mockRestore()
  vi.mocked(billingAttributionModule.resolveBillingAttribution).mockRestore()
  vi.mocked(billingAttributionModule.toBillingContext).mockRestore()
})

/**
 * Spy on the real documents/utils namespace instead of vi.mock: the shared
 * `@/lib/knowledge/embeddings` module may be cached bound to the real module,
 * so patching the namespace is the only wiring that always applies.
 */
const retrySpy = vi
  .spyOn(documentsUtilsModule, 'retryWithExponentialBackoff')
  .mockImplementation(((fn: () => unknown) => fn()) as never)

const BILLING_ATTRIBUTION_FIXTURE = {
  actorUserId: 'billing-user-1',
  billedAccountUserId: 'billing-user-1',
  billingEntity: { type: 'user', id: 'billing-user-1' },
  billingPeriod: {
    start: '2026-07-01T00:00:00.000Z',
    end: '2026-08-01T00:00:00.000Z',
  },
  organizationId: null,
  payerSubscription: null,
  workspaceId: 'workspace1',
} as never

function applyBillingSpies() {
  vi.spyOn(workspacesUtilsModule, 'getWorkspaceBilledAccountUserId').mockResolvedValue('user1')
  vi.spyOn(billingAttributionModule, 'assertBillingAttributionSnapshot').mockImplementation(
    ((value: unknown) => value) as never
  )
  vi.spyOn(billingAttributionModule, 'checkAttributedUsageLimits').mockResolvedValue({
    isExceeded: false,
    payerUsage: { currentUsage: 0, limit: 100 },
  } as never)
  vi.spyOn(billingAttributionModule, 'resolveBillingAttribution').mockResolvedValue(
    BILLING_ATTRIBUTION_FIXTURE
  )
  vi.spyOn(billingAttributionModule, 'toBillingContext').mockImplementation((() => ({
    billingEntity: { type: 'user', id: 'billing-user-1' },
    billingPeriod: {
      start: new Date('2026-07-01T00:00:00.000Z'),
      end: new Date('2026-08-01T00:00:00.000Z'),
    },
  })) as never)
}

/**
 * Billing helpers are spied on the real namespaces (not vi.mock'd) for the
 * same shared-consumer reason as the retry spy above.
 */
applyBillingSpies()

vi.mock('@/lib/knowledge/documents/document-processor', () => ({
  processDocument: vi.fn().mockResolvedValue({
    chunks: [
      {
        text: 'alpha',
        tokenCount: 1,
        metadata: { startIndex: 0, endIndex: 4 },
      },
      {
        text: 'beta',
        tokenCount: 1,
        metadata: { startIndex: 5, endIndex: 8 },
      },
    ],
    metadata: {
      filename: 'dummy',
      fileSize: 10,
      mimeType: 'text/plain',
      characterCount: 9,
      tokenCount: 3,
      chunkCount: 2,
      processingMethod: 'file-parser',
    },
  }),
}))

function createEmbeddingFetchMock() {
  return vi.fn().mockResolvedValue({
    ok: true,
    json: async () => ({
      data: [
        { embedding: [0.1, 0.2], index: 0 },
        { embedding: [0.3, 0.4], index: 1 },
      ],
      usage: { prompt_tokens: 2, total_tokens: 2 },
    }),
  })
}

vi.stubGlobal('fetch', createEmbeddingFetchMock())

import { processDocumentAsync } from '@/lib/knowledge/documents/service'
import { generateEmbeddings } from '@/lib/knowledge/embeddings'
import {
  checkChunkAccess,
  checkDocumentAccess,
  checkKnowledgeBaseAccess,
} from '@/app/api/knowledge/utils'

describe('Knowledge Utils', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    resetDbChainMock()
    // `unstubGlobals: true` removes the module-scope fetch stub after the
    // first test in the worker; re-stub it per test.
    vi.stubGlobal('fetch', createEmbeddingFetchMock())
    // Under `isolate: false` the shared `@/lib/knowledge/embeddings` module may
    // be cached bound to the REAL env module, so reset the real `env` object
    // per test instead of vi.mock'ing a file-local replacement that a cached
    // consumer would never see.
    for (const key of Object.keys(env)) {
      delete (env as Record<string, unknown>)[key]
    }
    Object.assign(env, { ...defaultMockEnv, OPENAI_API_KEY: 'test-key' })
    retrySpy.mockImplementation(((fn: () => unknown) => fn()) as never)
    applyBillingSpies()
  })

  describe('processDocumentAsync', () => {
    it('should insert embeddings before updating document counters', async () => {
      /** Context prefetch JOIN (document × knowledge_base × workspace). */
      queueTableRows(schemaMock.document, [
        {
          workspaceId: 'workspace1',
          knowledgeBaseUserId: 'user1',
          chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 200 },
          embeddingModel: 'text-embedding-3-small',
          billedAccountUserId: 'billing-user-1',
          uploadedBy: null,
        },
      ])
      /** Legacy untracked documents have exact-empty provenance. */
      queueTableRows(schemaMock.document, [
        {
          filename: 'file.txt',
          fileUrl: 'https://example.com/file.txt',
          secretProvenanceVersion: null,
        },
      ])
      /** In-transaction active-document recheck. */
      queueTableRows(schemaMock.document, [{ id: 'doc1' }])

      await processDocumentAsync(
        'kb1',
        'doc1',
        {
          filename: 'file.txt',
          fileUrl: 'https://example.com/file.txt',
          fileSize: 10,
          mimeType: 'text/plain',
        },
        {},
        {
          actorUserId: 'billing-user-1',
          billedAccountUserId: 'billing-user-1',
          billingEntity: { type: 'user', id: 'billing-user-1' },
          billingPeriod: {
            start: '2026-07-01T00:00:00.000Z',
            end: '2026-08-01T00:00:00.000Z',
          },
          organizationId: null,
          payerSubscription: null,
          workspaceId: 'workspace1',
        }
      )

      /**
       * Embeddings are inserted first, then the document counter update. The
       * status→'processing' update precedes both and a usage_log billing insert
       * (recordUsage) may trail after — assert relative order via the shared
       * spies' invocation order rather than exact call sequences.
       */
      const setPayloads = dbChainMockFns.set.mock.calls.map((call) => call[0])
      const completedIndex = setPayloads.findIndex((p) => p?.processingStatus === 'completed')
      expect(setPayloads[completedIndex]).toMatchObject({
        processingStatus: 'completed',
        chunkCount: 2,
      })

      expect(dbChainMockFns.values.mock.calls[0][0]).toHaveLength(2)
      expect(dbChainMockFns.values.mock.invocationCallOrder[0]).toBeLessThan(
        dbChainMockFns.set.mock.invocationCallOrder[completedIndex]
      )
    })
  })

  describe('checkKnowledgeBaseAccess', () => {
    it('should return success for owner', async () => {
      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
      const result = await checkKnowledgeBaseAccess('kb1', 'user1')

      expect(result.hasAccess).toBe(true)
    })

    it('should return notFound when knowledge base is missing', async () => {
      const result = await checkKnowledgeBaseAccess('missing', 'user1')

      expect(result.hasAccess).toBe(false)
      expect('notFound' in result && result.notFound).toBe(true)
    })
  })

  describe('checkDocumentAccess', () => {
    it('should return unauthorized when user mismatch', async () => {
      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'owner' }])
      const result = await checkDocumentAccess('kb1', 'doc1', 'intruder')

      expect(result.hasAccess).toBe(false)
      if ('reason' in result) {
        expect(result.reason).toBe('Unauthorized knowledge base access')
      }
    })
  })

  describe('checkChunkAccess', () => {
    it('should fail when document is not completed', async () => {
      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
      queueTableRows(schemaMock.document, [
        { id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'processing' },
      ])

      const result = await checkChunkAccess('kb1', 'doc1', 'chunk1', 'user1')

      expect(result.hasAccess).toBe(false)
      if ('reason' in result) {
        expect(result.reason).toContain('Document is not ready')
      }
    })

    it('should return success for valid access', async () => {
      queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb1', userId: 'user1' }])
      queueTableRows(schemaMock.document, [
        { id: 'doc1', knowledgeBaseId: 'kb1', processingStatus: 'completed' },
      ])
      queueTableRows(schemaMock.embedding, [{ id: 'chunk1', documentId: 'doc1' }])

      const result = await checkChunkAccess('kb1', 'doc1', 'chunk1', 'user1')

      expect(result.hasAccess).toBe(true)
      if ('chunk' in result) {
        expect(result.chunk.id).toBe('chunk1')
      }
    })
  })

  describe('generateEmbeddings', () => {
    it('should return same length as input', async () => {
      const result = await generateEmbeddings(['a', 'b'])

      expect(result.embeddings.length).toBe(2)
    })

    it('should use Azure OpenAI when Azure config is provided', async () => {
      const { env } = await import('@/lib/core/config/env')
      Object.keys(env).forEach((key) => delete (env as any)[key])
      Object.assign(env, {
        AZURE_OPENAI_API_KEY: 'test-azure-key',
        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
        OPENAI_API_KEY: 'test-openai-key',
      })

      const fetchSpy = vi.mocked(fetch)
      fetchSpy.mockResolvedValueOnce({
        ok: true,
        json: async () => ({
          data: [{ embedding: [0.1, 0.2], index: 0 }],
          usage: { prompt_tokens: 1, total_tokens: 1 },
        }),
      } as any)

      await generateEmbeddings(['test text'])

      expect(fetchSpy).toHaveBeenCalledWith(
        'https://test.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-12-01-preview',
        expect.objectContaining({
          headers: expect.objectContaining({
            'api-key': 'test-azure-key',
          }),
        })
      )

      Object.keys(env).forEach((key) => delete (env as any)[key])
    })

    it('should fallback to OpenAI when no Azure config provided', async () => {
      const { env } = await import('@/lib/core/config/env')
      Object.keys(env).forEach((key) => delete (env as any)[key])
      Object.assign(env, {
        OPENAI_API_KEY: 'test-openai-key',
      })

      const fetchSpy = vi.mocked(fetch)
      fetchSpy.mockResolvedValueOnce({
        ok: true,
        json: async () => ({
          data: [{ embedding: [0.1, 0.2], index: 0 }],
          usage: { prompt_tokens: 1, total_tokens: 1 },
        }),
      } as any)

      await generateEmbeddings(['test text'])

      expect(fetchSpy).toHaveBeenCalledWith(
        'https://api.openai.com/v1/embeddings',
        expect.objectContaining({
          headers: expect.objectContaining({
            Authorization: 'Bearer test-openai-key',
          }),
        })
      )

      Object.keys(env).forEach((key) => delete (env as any)[key])
    })

    it('should throw error when no API configuration provided', async () => {
      const { env } = await import('@/lib/core/config/env')
      Object.keys(env).forEach((key) => delete (env as any)[key])
      // The env object lazily reads process.env, so a developer's local .env
      // keys survive the deletion above — stub the direct key empty and fail
      // the hosted rotation fallback for hermeticity on any machine.
      vi.stubEnv('OPENAI_API_KEY', '')
      const apiKeysModule = await import('@/lib/core/config/api-keys')
      const rotationSpy = vi.spyOn(apiKeysModule, 'getRotatingApiKey').mockImplementation(() => {
        throw new Error('No rotation keys configured')
      })

      try {
        await expect(generateEmbeddings(['test text'])).rejects.toThrow(
          'OPENAI_API_KEY is not configured'
        )
      } finally {
        rotationSpy.mockRestore()
        vi.unstubAllEnvs()
      }
    })
  })
})