ai-drama-platform/scripts/smoke-worker.mjs

import assert from "node:assert/strict";
import { createServer } from "node:http";
import { dbRun, withTransaction } from "../server/db.mjs";

const api = process.env.AI_DRAMA_API_BASE || "http://127.0.0.1:8787";
const runnerPort = Number(process.env.AI_DRAMA_SMOKE_RUNNER_PORT || 8791);
const scope = {
  "x-organization-id": "org-studio-lab",
  "x-workspace-id": "ws-local-aidrama",
  "x-project-id": "thunder-mouth"
};

async function request(path, options = {}) {
  const response = await fetch(`${api}${path}`, {
    ...options,
    headers: { "content-type": "application/json", ...(options.headers || {}) }
  });
  const payload = await response.json().catch(() => ({}));
  return { response, payload };
}

const runnerRequests = [];
const runner = createServer(async (req, res) => {
  const chunks = [];
  for await (const chunk of req) chunks.push(chunk);
  const body = Buffer.concat(chunks).toString("utf8");
  runnerRequests.push({ method: req.method, url: req.url, body: body ? JSON.parse(body) : null });
  res.writeHead(200, { "content-type": "application/json" });
  if (req.url === "/v1/images/generations") {
    res.end(JSON.stringify({ created: Date.now(), data: [{ b64_json: "smoke-image" }] }));
    return;
  }
  res.end(JSON.stringify({ outputPath: "storage/jobs/smoke-custom/output.json", data: [{ id: "single-output" }] }));
});

await new Promise((resolve) => runner.listen(runnerPort, "127.0.0.1", resolve));
let customModelId = "";
let openAiModelId = "";
let jobIds = [];

try {
  const login = await request("/api/auth/login", {
    method: "POST",
    body: JSON.stringify({ email: "producer@local.test", password: "Demo@123456" })
  });
  assert.equal(login.response.ok, true, "worker smoke login failed");
  const headers = { authorization: `Bearer ${login.payload.session.token}`, ...scope };

  const customModel = await request("/api/platform/models/register", {
    method: "POST",
    headers,
    body: JSON.stringify({
      label: `Smoke Custom Runner ${Date.now()}`,
      endpoint: `http://127.0.0.1:${runnerPort}/custom/generate`,
      kind: "http-json",
      capability: ["text-to-image", "single-frame"],
      costMode: "local",
      protocol: { healthRoute: "custom/health" }
    })
  });
  assert.equal(customModel.response.status, 201, "custom runner registration failed");
  customModelId = customModel.payload.model.id;
  await request(`/api/platform/models/${encodeURIComponent(customModelId)}`, {
    method: "PATCH",
    headers,
    body: JSON.stringify({ status: "ready" })
  });

  const openAiModel = await request("/api/platform/models/register", {
    method: "POST",
    headers,
    body: JSON.stringify({
      label: `Smoke OpenAI Runner ${Date.now()}`,
      endpoint: `http://127.0.0.1:${runnerPort}/v1`,
      kind: "openai-compatible",
      capability: ["text-to-image", "single-frame"],
      costMode: "local",
      protocol: { models: { image: "smoke-image" }, routes: { image: "images/generations" } }
    })
  });
  assert.equal(openAiModel.response.status, 201, "OpenAI-compatible runner registration failed");
  openAiModelId = openAiModel.payload.model.id;
  await request(`/api/platform/models/${encodeURIComponent(openAiModelId)}`, {
    method: "PATCH",
    headers,
    body: JSON.stringify({ status: "ready" })
  });

  const customJob = await request("/api/jobs", {
    method: "POST",
    headers,
    body: JSON.stringify({ adapter: customModelId, kind: "自定义单画面关键帧", shotId: "shot-01", output: "storage/jobs/smoke-custom/output.json" })
  });
  assert.equal(customJob.response.status, 201, "custom job creation failed");
  assert.equal(customJob.payload.job.status, "queued", "custom job should be queued for local Worker");

  const openAiJob = await request("/api/jobs", {
    method: "POST",
    headers,
    body: JSON.stringify({ adapter: openAiModelId, kind: "OpenAI-compatible 单画面关键帧", shotId: "shot-01", output: "storage/jobs/smoke-openai/output.png" })
  });
  assert.equal(openAiJob.response.status, 201, "OpenAI-compatible job creation failed");
  assert.equal(openAiJob.payload.job.status, "queued", "OpenAI-compatible job should be queued");
  jobIds = [customJob.payload.job.id, openAiJob.payload.job.id];

  const deadline = Date.now() + 12_000;
  const completed = new Map();
  while (Date.now() < deadline && completed.size < jobIds.length) {
    for (const jobId of jobIds) {
      const result = await request(`/api/jobs/${encodeURIComponent(jobId)}`, { headers });
      assert.equal(result.response.ok, true, `job ${jobId} detail request failed`);
      if (result.payload.job.status === "completed") completed.set(jobId, result.payload.job);
      if (result.payload.job.status === "failed") throw new Error(`${jobId} failed: ${result.payload.job.errorMessage}`);
    }
    if (completed.size < jobIds.length) await new Promise((resolve) => setTimeout(resolve, 400));
  }
  assert.equal(completed.size, jobIds.length, "local Worker did not complete all smoke jobs in time");
  assert.ok(completed.get(customJob.payload.job.id).result?.data?.length === 1, "custom protocol must preserve exactly one image output");
  assert.ok(completed.get(openAiJob.payload.job.id).result?.data?.length === 1, "OpenAI-compatible protocol must preserve exactly one image output");
  assert.ok(runnerRequests.some((item) => item.url === "/v1/images/generations"), "OpenAI-compatible route was not called");
  assert.ok(runnerRequests.some((item) => item.url === "/custom/generate"), "custom JSON route was not called");
  assert.ok(completed.get(customJob.payload.job.id).leased_by == null, "completed job lease must be released");
  console.log(`worker smoke passed: ${jobIds.length} jobs completed through local protocols`);
} finally {
  await new Promise((resolve) => runner.close(resolve));
  if (customModelId || openAiModelId || jobIds.length) {
    withTransaction(() => {
      for (const jobId of jobIds) {
        dbRun("DELETE FROM media_artifacts WHERE job_id = ?", [jobId]);
        dbRun("DELETE FROM job_dependencies WHERE job_id = ? OR depends_on_job_id = ?", [jobId, jobId]);
        dbRun("DELETE FROM job_attempts WHERE job_id = ?", [jobId]);
        dbRun("DELETE FROM usage_events WHERE metadata_json LIKE ?", [`%${jobId}%`]);
        dbRun("DELETE FROM audit_logs WHERE target_id = ?", [jobId]);
        dbRun("DELETE FROM generation_jobs WHERE id = ?", [jobId]);
      }
      for (const modelId of [customModelId, openAiModelId].filter(Boolean)) {
        dbRun("DELETE FROM model_connectors WHERE id = ?", [modelId]);
        dbRun("DELETE FROM audit_logs WHERE target_id = ?", [modelId]);
      }
    });
  }
}