easyai-ai-gateway/scripts/cluster/cross-node-file-e2e.mjs

#!/usr/bin/env node

import { createHash, randomUUID } from 'node:crypto'
import { deflateSync } from 'node:zlib'

function assert(condition, message) {
  if (!condition) throw new Error(message)
}

function crc32(buffer) {
  let value = 0xffffffff
  for (const byte of buffer) {
    value ^= byte
    for (let bit = 0; bit < 8; bit += 1) {
      value = (value >>> 1) ^ (value & 1 ? 0xedb88320 : 0)
    }
  }
  return (value ^ 0xffffffff) >>> 0
}

function pngChunk(type, data) {
  const typeBuffer = Buffer.from(type)
  const length = Buffer.alloc(4)
  length.writeUInt32BE(data.length)
  const checksum = Buffer.alloc(4)
  checksum.writeUInt32BE(crc32(Buffer.concat([typeBuffer, data])))
  return Buffer.concat([length, typeBuffer, data, checksum])
}

function createFixturePNG(width = 256, height = 256) {
  const header = Buffer.alloc(13)
  header.writeUInt32BE(width, 0)
  header.writeUInt32BE(height, 4)
  header[8] = 8
  header[9] = 2
  const rows = []
  for (let y = 0; y < height; y += 1) {
    const row = Buffer.alloc(1 + width * 3)
    for (let x = 0; x < width; x += 1) {
      row[1 + x * 3] = Math.round((x / width) * 255)
      row[2 + x * 3] = Math.round((y / height) * 255)
      row[3 + x * 3] = 180
    }
    rows.push(row)
  }
  return Buffer.concat([
    Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]),
    pngChunk('IHDR', header),
    pngChunk('IDAT', deflateSync(Buffer.concat(rows))),
    pngChunk('IEND', Buffer.alloc(0)),
  ])
}

const baseURL = process.env.AI_GATEWAY_E2E_BASE_URL || 'http://10.77.0.1:18089'
const account = Buffer.from(process.env.E2E_ACCOUNT_B64 || '', 'base64').toString('utf8')
const password = Buffer.from(process.env.E2E_PASSWORD_B64 || '', 'base64').toString('utf8')
const model = process.env.AI_GATEWAY_E2E_IMAGE_MODEL || 'openai:gpt-image-1'
const imageSize = process.env.AI_GATEWAY_E2E_IMAGE_SIZE || '2048x2048'
assert(account && password, 'E2E login credentials are unavailable')
assert(/^\d+x\d+$/.test(imageSize), 'E2E image size is invalid')
const [imageWidth, imageHeight] = imageSize.split('x').map(Number)
const imageResolution =
  process.env.AI_GATEWAY_E2E_IMAGE_RESOLUTION ||
  (imageWidth === imageHeight && imageWidth % 1024 === 0 ? `${imageWidth / 1024}K` : '')
const imageAspectRatio =
  process.env.AI_GATEWAY_E2E_IMAGE_ASPECT_RATIO ||
  `${imageWidth / greatestCommonDivisor(imageWidth, imageHeight)}:${imageHeight / greatestCommonDivisor(imageWidth, imageHeight)}`
assert(imageResolution, 'E2E image resolution is unavailable')

function greatestCommonDivisor(left, right) {
  while (right !== 0) [left, right] = [right, left % right]
  return left
}

async function request(path, { method = 'GET', token, body, form } = {}) {
  const headers = { Host: 'ai.51easyai.com' }
  if (token) headers.Authorization = `Bearer ${token}`
  if (body !== undefined) headers['Content-Type'] = 'application/json'
  if (form) headers['X-Async'] = 'true'
  const response = await fetch(`${baseURL}${path}`, {
    method,
    headers,
    body: form || (body === undefined ? undefined : JSON.stringify(body)),
    signal: AbortSignal.timeout(3_600_000),
  })
  const text = await response.text()
  let payload
  try {
    payload = text ? JSON.parse(text) : {}
  } catch {
    payload = { raw: text.slice(0, 500) }
  }
  if (!response.ok) {
    throw new Error(`${method} ${path} returned HTTP ${response.status}: ${JSON.stringify(payload)}`)
  }
  return { response, payload }
}

function collectHTTPSURLs(value, output = []) {
  if (typeof value === 'string' && value.startsWith('https://')) output.push(value)
  else if (Array.isArray(value)) value.forEach((item) => collectHTTPSURLs(item, output))
  else if (value && typeof value === 'object') Object.values(value).forEach((item) => collectHTTPSURLs(item, output))
  return output
}

const login = await request('/api/v1/auth/login', {
  method: 'POST',
  body: { account, password },
})
const jwt = login.payload.accessToken
assert(typeof jwt === 'string' && jwt.length > 20, 'login did not return an access token')

const existingKeys = await request('/api/v1/api-keys', { token: jwt })
for (const item of existingKeys.payload.items || []) {
  if (typeof item?.id === 'string' && item?.name?.startsWith('cluster-cross-node-')) {
    await request(`/api/v1/api-keys/${item.id}`, { method: 'DELETE', token: jwt })
  }
}

const keyName = `cluster-cross-node-${Date.now()}`
const keyResult = await request('/api/v1/api-keys', {
  method: 'POST',
  token: jwt,
  body: { name: keyName },
})
const apiKey = keyResult.payload.secret
const apiKeyId = keyResult.payload.apiKey?.id
assert(typeof apiKey === 'string' && apiKey.startsWith('sk-gw-'), 'API key creation failed')
assert(typeof apiKeyId === 'string' && apiKeyId.length > 0, 'API key ID is unavailable')

let result
let taskId
let taskSucceeded = false
try {
  const assignableModels = await request('/api/v1/api-keys/assignable-models', { token: jwt })
  const targetModels = (assignableModels.payload.items || []).filter(
    (item) =>
      [item?.modelName, item?.modelAlias].includes(model) &&
      Array.isArray(item?.modelType) &&
      item.modelType.includes('image_edit'),
  )
  assert(targetModels.length > 0, `no assignable image_edit platform model found for ${model}`)
  const resourceMap = new Map()
  for (const item of targetModels) {
    for (const [resourceType, resourceId] of [
      ['platform', item.platformId],
      ['platform_model', item.id],
      ['base_model', item.baseModelId],
    ]) {
      if (typeof resourceId === 'string' && resourceId) {
        resourceMap.set(`${resourceType}:${resourceId}`, {
          resourceType,
          resourceId,
          priority: 100,
          minPermissionLevel: 0,
          status: 'active',
        })
      }
    }
  }
  await request('/api/v1/api-keys/access-rules/batch', {
    method: 'POST',
    token: jwt,
    body: {
      subjectType: 'api_key',
      subjectId: apiKeyId,
      effect: 'allow',
      upsertResources: [...resourceMap.values()],
      deleteResources: [],
    },
  })

  const fixture = createFixturePNG()
  const sourceSHA256 = createHash('sha256').update(fixture).digest('hex')
  const form = new FormData()
  form.append('model', model)
  form.append('prompt', `Cross-node HA acceptance ${randomUUID()}: add a thin white border`)
  form.append('size', imageSize)
  form.append('width', String(imageWidth))
  form.append('height', String(imageHeight))
  form.append('resolution', imageResolution)
  form.append('aspect_ratio', imageAspectRatio)
  form.append('images', new Blob([fixture], { type: 'image/png' }), 'cluster-e2e.png')
  const accepted = await request('/api/v1/images/edits', {
    method: 'POST',
    token: apiKey,
    form,
  })
  assert(accepted.response.status === 202, `image edit returned ${accepted.response.status}, expected 202`)
  taskId = accepted.payload.taskId || accepted.payload.task_id || accepted.payload.task?.id
  assert(/^[0-9a-f-]{36}$/.test(taskId || ''), 'image edit response has no task ID')

  let task
  const deadline = Date.now() + 20 * 60_000
  while (Date.now() < deadline) {
    const taskResult = await request(`/api/v1/tasks/${taskId}`, { token: apiKey })
    task = taskResult.payload.task || taskResult.payload
    if (task.status === 'succeeded') break
    if (['failed', 'cancelled'].includes(task.status)) {
      throw new Error(`task ${taskId} ended as ${task.status}: ${task.errorMessage || task.error || ''}`)
    }
    await new Promise((resolve) => setTimeout(resolve, 3000))
  }
  assert(task?.status === 'succeeded', `task ${taskId} did not succeed before timeout`)
  const resultURLs = [...new Set(collectHTTPSURLs(task.result))]
  assert(resultURLs.length > 0, 'task result contains no shared HTTPS URL')
  taskSucceeded = true
  result = { taskId, sourceSHA256, resultURL: resultURLs[0], model }
} finally {
  if (taskId && !taskSucceeded) {
    await request(`/api/v1/tasks/${taskId}/cancel`, { method: 'POST', token: apiKey }).catch(() => {})
  }
  if (apiKeyId) {
    await request(`/api/v1/api-keys/${apiKeyId}`, { method: 'DELETE', token: jwt }).catch(() => {})
  }
}

console.log(JSON.stringify(result))