dify/packages/contracts/scripts/generate-knowledge-fs-contract.mjs

import { execFileSync } from 'node:child_process'
import { createHash } from 'node:crypto'
import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { dirname, join, resolve } from 'node:path'
import { fileURLToPath } from 'node:url'
import { getStreamingOperationIds } from './knowledge-fs-contract-utils.mjs'

const packageRoot = resolve(dirname(fileURLToPath(import.meta.url)), '..')
const workspaceRoot = resolve(packageRoot, '../..')
const repository = resolve(
  process.env.KNOWLEDGE_FS_REPO ?? resolve(workspaceRoot, '../knowledge-fs'),
)
const temporaryDirectory = await mkdtemp(join(tmpdir(), 'dify-knowledge-fs-types-'))

try {
  const openapiPath = join(temporaryDirectory, 'knowledge-fs.console.json')
  run(
    'uv',
    [
      'run',
      '--project',
      resolve(workspaceRoot, 'api'),
      resolve(workspaceRoot, 'api/dev/generate_knowledge_fs_contract.py'),
      '--repository',
      repository,
      '--check',
      '--output-openapi',
      openapiPath,
    ],
    workspaceRoot,
  )
  run('pnpm', ['exec', 'openapi-ts', '-f', 'openapi-ts.knowledge-fs.config.ts'], packageRoot, {
    KNOWLEDGE_FS_OPENAPI: openapiPath,
  })
  await patchStreamingContracts(openapiPath)
  run('pnpm', ['exec', 'vp', 'fmt', 'generated/knowledge-fs'], packageRoot)
  await writeContractMetadata(openapiPath, await generatedArtifactSha256())
  run('pnpm', ['exec', 'vp', 'fmt', 'generated/knowledge-fs/metadata.gen.ts'], packageRoot)
} finally {
  await rm(temporaryDirectory, { force: true, recursive: true })
}

async function patchStreamingContracts(openapiPath) {
  const document = JSON.parse(await readFile(openapiPath, 'utf8'))
  const streamingOperationIds = getStreamingOperationIds(document)
  if (streamingOperationIds.length === 0) return

  const outputPath = join(packageRoot, 'generated/knowledge-fs/orpc.gen.ts')
  let source = await readFile(outputPath, 'utf8')
  source = replaceOnce(
    source,
    "import { oc } from '@orpc/contract'",
    "import { eventIterator, oc } from '@orpc/contract'",
  )

  for (const operationId of streamingOperationIds) {
    const responseSchema = `z${capitalize(operationId)}Response`
    source = replaceOnce(
      source,
      `.output(${responseSchema})`,
      `.output(eventIterator(${responseSchema}))`,
    )
  }

  await writeFile(outputPath, source)
}

function capitalize(value) {
  return value.charAt(0).toUpperCase() + value.slice(1)
}

function replaceOnce(source, target, replacement) {
  const firstIndex = source.indexOf(target)
  if (firstIndex === -1 || source.indexOf(target, firstIndex + target.length) !== -1)
    throw new Error(`Expected exactly one generated occurrence of ${target}`)

  return source.slice(0, firstIndex) + replacement + source.slice(firstIndex + target.length)
}

async function generatedArtifactSha256() {
  const generatedDirectory = join(packageRoot, 'generated/knowledge-fs')
  const fileNames = (await readdir(generatedDirectory))
    .filter((fileName) => fileName.endsWith('.gen.ts') && fileName !== 'metadata.gen.ts')
    .sort()

  return Object.fromEntries(
    await Promise.all(
      fileNames.map(async (fileName) => [
        fileName,
        createHash('sha256')
          .update(await readFile(join(generatedDirectory, fileName)))
          .digest('hex'),
      ]),
    ),
  )
}

async function writeContractMetadata(openapiPath, artifactSha256) {
  const document = JSON.parse(await readFile(openapiPath, 'utf8'))
  const source = [
    '// This file is auto-generated by scripts/generate-knowledge-fs-contract.mjs.',
    '// Do not edit it manually.',
    '',
    `export const knowledgeFsSourceOpenapiSha256 = ${JSON.stringify(document['x-dify-source-openapi-sha256'])}`,
    `export const knowledgeFsConsoleDeclarationsSha256 = ${JSON.stringify(document['x-dify-console-declarations-sha256'])}`,
    `export const knowledgeFsGeneratedArtifactSha256 = ${JSON.stringify(artifactSha256, null, 2)} as const`,
    '',
  ].join('\n')
  await writeFile(join(packageRoot, 'generated/knowledge-fs/metadata.gen.ts'), source)
}

function run(command, args, cwd, extraEnv = {}) {
  execFileSync(command, args, {
    cwd,
    env: { ...process.env, ...extraEnv },
    stdio: 'inherit',
  })
}