botpress--botpress/packages/zai/e2e/client.ts

import { Client } from '@botpress/client'
import { Cognitive } from '@botpress/cognitive'
import { diffLines } from 'diff'
import fs from 'node:fs'
import path from 'node:path'
import { expect } from 'vitest'

function stringifyWithSortedKeys(obj: any, space?: number): string {
  function sortKeys(input: any): any {
    if (Array.isArray(input)) {
      return input.map(sortKeys)
    } else if (input && typeof input === 'object' && input.constructor === Object) {
      return Object.keys(input)
        .sort()
        .reduce(
          (acc, key) => {
            acc[key] = sortKeys(input[key])
            return acc
          },
          {} as Record<string, any>
        )
    } else {
      return input
    }
  }

  return JSON.stringify(sortKeys(obj), null, space)
}

function readJSONL<T>(filePath: string, keyProperty: keyof T): Map<string, T> {
  const lines = fs.readFileSync(filePath, 'utf-8').split(/\r?\n/).filter(Boolean)

  const map = new Map<string, T>()

  for (const line of lines) {
    try {
      const obj = JSON.parse(line) as T
      const value = obj[keyProperty]
      if (value === undefined || value === null) {
        continue
      }
      const key = String(value)
      map.set(key, obj)
    } catch {}
  }

  return map
}

type CacheEntry = { key: string; value: any; test?: string; input: string }

const cache: Map<string, CacheEntry> = readJSONL(path.resolve(__dirname, './data/cache.jsonl'), 'key')
const cacheByTest: Map<string, CacheEntry> = readJSONL(path.resolve(__dirname, './data/cache.jsonl'), 'test')

class CachedClient extends Client {
  #client: Client
  #callsByTest: Record<string, number> = {}

  public constructor(options: ConstructorParameters<typeof Client>[0]) {
    super(options)
    this.#client = new Client(options)
  }

  public callAction = async (...args: Parameters<Client['callAction']>) => {
    const currentTestName = expect.getState().currentTestName ?? 'default'
    this.#callsByTest[currentTestName] ||= 0
    this.#callsByTest[currentTestName]++

    const testKey = `${currentTestName}-${this.#callsByTest[currentTestName]}`

    const key = fastHash(stringifyWithSortedKeys(args))
    const cached = cache.get(key)

    if (cached) {
      return cached.value
    }

    if (process.env.CI) {
      const previous = cacheByTest.get(testKey)
      if (previous) {
        console.info(`Cache miss for ${key} in test ${testKey}`)
        console.info(
          diffLines(
            JSON.stringify(JSON.parse(previous.input), null, 2),
            JSON.stringify(JSON.parse(stringifyWithSortedKeys(args)), null, 2)
          )
        )
      }

      throw new Error(
        `Missing cached Botpress action response for ${key} in ${testKey}. Run pnpm test:e2e:update to refresh packages/zai/e2e/data/cache.jsonl.`
      )
    }

    const response = await this.#client.callAction(...args)
    cache.set(key, { key, value: response, test: testKey, input: stringifyWithSortedKeys(args) })

    fs.appendFileSync(
      path.resolve(__dirname, './data/cache.jsonl'),
      JSON.stringify({
        test: testKey,
        key,
        input: stringifyWithSortedKeys(args),
        value: response,
      }) + '\n'
    )

    return response
  }

  public clone() {
    return this
  }
}

export const getCognitiveClient = () => {
  const cognitive = new Cognitive({
    client: new Client({
      apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
      botId: process.env.CLOUD_BOT_ID,
      token: process.env.CLOUD_PAT,
    }),
    __experimental_beta: true,
  })
  return cognitive
}

export const getCachedCognitiveClient = () => {
  const cognitive = new Cognitive({
    client: new CachedClient({
      // Using mock server (MSW) - actual URL doesn't matter as requests are intercepted
      apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
      botId: process.env.CLOUD_BOT_ID,
      token: process.env.CLOUD_PAT,
    }),
    __experimental_beta: true,
  })
  return cognitive
}

function fastHash(str: string): string {
  let hash = 0
  for (let i = 0; i < str.length; i++) {
    hash = (hash << 5) - hash + str.charCodeAt(i)
    hash |= 0 // Convert to 32bit integer
  }
  return (hash >>> 0).toString(16) // Convert to unsigned and then to hex
}