stringadmin
1083 lines
36 KiB
JavaScript
1083 lines
36 KiB
JavaScript
"use strict";
|
||
|
||
const crypto = require("node:crypto");
|
||
const { EventEmitter } = require("node:events");
|
||
const { pool } = require("./db");
|
||
const { refundTaskBillingOnFailure } = require("./billing");
|
||
const { putObject, isOssConfigured } = require("./ossClient");
|
||
|
||
const taskEvents = new EventEmitter();
|
||
taskEvents.setMaxListeners(200);
|
||
|
||
const activePollers = new Map();
|
||
const POLL_INTERVAL_MS = 3000;
|
||
const MAX_POLL_ATTEMPTS = 120;
|
||
const GRS_IMAGE_MAX_POLL_ATTEMPTS = Number(process.env.GRSAI_IMAGE_MAX_POLL_ATTEMPTS || 60);
|
||
const TASK_EVENT_CHANNEL = "generation_task_events";
|
||
const TASK_EVENT_ORIGIN = `${process.pid}-${crypto.randomUUID()}`;
|
||
const POLLER_OWNER_ID = `${process.pid}-${crypto.randomUUID()}`;
|
||
const POLLER_OWNER_STALE_MS = Number(process.env.TASK_POLLER_OWNER_STALE_MS || 20_000);
|
||
const POLLER_RECOVERY_INTERVAL_MS = Number(process.env.TASK_POLLER_RECOVERY_INTERVAL_MS || 30_000);
|
||
let taskEventListenerClient = null;
|
||
let taskEventListenerStarting = null;
|
||
let pollerStoreReady = null;
|
||
let pollerRecoveryTimer = null;
|
||
|
||
function normalizeTaskProgress(value) {
|
||
const numeric = Number(value);
|
||
if (!Number.isFinite(numeric)) return undefined;
|
||
return Math.max(0, Math.min(100, Math.round(numeric)));
|
||
}
|
||
|
||
function formatTaskEvent(row) {
|
||
return {
|
||
taskId: row.id,
|
||
status: row.status,
|
||
progress: row.progress,
|
||
resultUrl: row.result_url || null,
|
||
error: row.error || null,
|
||
};
|
||
}
|
||
|
||
function emitTaskEvent(event) {
|
||
if (!event?.taskId) return;
|
||
taskEvents.emit(`task:${event.taskId}`, event);
|
||
}
|
||
|
||
async function publishTaskEvent(event) {
|
||
if (!event?.taskId) return;
|
||
emitTaskEvent(event);
|
||
try {
|
||
await pool.query("SELECT pg_notify($1, $2)", [
|
||
TASK_EVENT_CHANNEL,
|
||
JSON.stringify({ origin: TASK_EVENT_ORIGIN, event }),
|
||
]);
|
||
} catch (err) {
|
||
console.error(`[aiTaskWorker] task event publish failed for task ${event.taskId}:`, err.message);
|
||
}
|
||
}
|
||
|
||
function serializeProviderConfig(providerConfig) {
|
||
if (!providerConfig || typeof providerConfig !== "object") return {};
|
||
const allowedKeys = [
|
||
"provider",
|
||
"transport",
|
||
"protocol",
|
||
"baseUrl",
|
||
"endpoint",
|
||
"resultEndpoint",
|
||
"model",
|
||
"requestedModel",
|
||
];
|
||
const result = {};
|
||
for (const key of allowedKeys) {
|
||
if (providerConfig[key] !== undefined) result[key] = providerConfig[key];
|
||
}
|
||
return result;
|
||
}
|
||
|
||
function parseProviderConfig(value) {
|
||
if (!value) return {};
|
||
if (typeof value === "object") return value;
|
||
try {
|
||
const parsed = JSON.parse(value);
|
||
return parsed && typeof parsed === "object" ? parsed : {};
|
||
} catch {
|
||
return {};
|
||
}
|
||
}
|
||
|
||
async function ensureTaskPollerStore() {
|
||
if (pollerStoreReady) return pollerStoreReady;
|
||
pollerStoreReady = pool.query(`
|
||
CREATE TABLE IF NOT EXISTS generation_task_pollers (
|
||
task_id INTEGER PRIMARY KEY REFERENCES generation_tasks(id) ON DELETE CASCADE,
|
||
provider_task_id TEXT NOT NULL,
|
||
task_type TEXT NOT NULL,
|
||
provider_config_json TEXT NOT NULL,
|
||
lease_token TEXT,
|
||
owner_id TEXT,
|
||
owner_heartbeat_at TIMESTAMPTZ,
|
||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_generation_task_pollers_owner
|
||
ON generation_task_pollers(owner_heartbeat_at);
|
||
`).catch((err) => {
|
||
pollerStoreReady = null;
|
||
throw err;
|
||
});
|
||
return pollerStoreReady;
|
||
}
|
||
|
||
async function persistPollerState(taskDbId, { providerTaskId, type, providerConfig, leaseToken }) {
|
||
await ensureTaskPollerStore();
|
||
await pool.query(
|
||
`
|
||
INSERT INTO generation_task_pollers (
|
||
task_id, provider_task_id, task_type, provider_config_json, lease_token,
|
||
owner_id, owner_heartbeat_at, updated_at
|
||
)
|
||
VALUES ($1, $2, $3, $4, $5, $6, NOW(), NOW())
|
||
ON CONFLICT (task_id) DO UPDATE SET
|
||
provider_task_id = EXCLUDED.provider_task_id,
|
||
task_type = EXCLUDED.task_type,
|
||
provider_config_json = EXCLUDED.provider_config_json,
|
||
lease_token = EXCLUDED.lease_token,
|
||
owner_id = EXCLUDED.owner_id,
|
||
owner_heartbeat_at = NOW(),
|
||
updated_at = NOW()
|
||
`,
|
||
[
|
||
taskDbId,
|
||
providerTaskId,
|
||
type,
|
||
JSON.stringify(serializeProviderConfig(providerConfig)),
|
||
leaseToken || null,
|
||
POLLER_OWNER_ID,
|
||
],
|
||
);
|
||
}
|
||
|
||
async function refreshPollerHeartbeat(taskDbId) {
|
||
await ensureTaskPollerStore();
|
||
await pool.query(
|
||
"UPDATE generation_task_pollers SET owner_id = $1, owner_heartbeat_at = NOW(), updated_at = NOW() WHERE task_id = $2",
|
||
[POLLER_OWNER_ID, taskDbId],
|
||
);
|
||
}
|
||
|
||
async function clearPollerState(taskDbId) {
|
||
await ensureTaskPollerStore();
|
||
await pool.query("DELETE FROM generation_task_pollers WHERE task_id = $1", [taskDbId]);
|
||
}
|
||
|
||
async function getLeaseKey(leaseToken) {
|
||
if (!leaseToken) return null;
|
||
const { rows } = await pool.query(
|
||
`
|
||
SELECT k.api_key
|
||
FROM key_leases l
|
||
JOIN api_keys k ON k.id = l.key_id
|
||
WHERE l.lease_token = $1
|
||
AND l.released_at IS NULL
|
||
AND k.enabled = 1
|
||
LIMIT 1
|
||
`,
|
||
[leaseToken],
|
||
);
|
||
const apiKey = rows[0]?.api_key;
|
||
return apiKey === "pool-slot" ? "" : apiKey || null;
|
||
}
|
||
|
||
async function claimPoller(taskId) {
|
||
await ensureTaskPollerStore();
|
||
const staleInterval = `${Math.max(5, Math.ceil(POLLER_OWNER_STALE_MS / 1000))} seconds`;
|
||
const { rows } = await pool.query(
|
||
`
|
||
UPDATE generation_task_pollers
|
||
SET owner_id = $1, owner_heartbeat_at = NOW(), updated_at = NOW()
|
||
WHERE task_id = $2
|
||
AND (
|
||
owner_heartbeat_at IS NULL
|
||
OR owner_heartbeat_at < NOW() - ($3::text)::interval
|
||
)
|
||
RETURNING *
|
||
`,
|
||
[POLLER_OWNER_ID, taskId, staleInterval],
|
||
);
|
||
return rows[0] || null;
|
||
}
|
||
|
||
async function createTaskLifecycleNotification(task) {
|
||
if (!task || !task.user_id || !task.id) return;
|
||
|
||
const isCompleted = task.status === "completed";
|
||
const isFailed = task.status === "failed";
|
||
if (!isCompleted && !isFailed) return;
|
||
|
||
const typeLabel = task.type === "video" ? "视频" : "图像";
|
||
const noticeType = isCompleted ? "task_completed" : "task_failed";
|
||
const title = isCompleted ? `${typeLabel}生成已完成` : `${typeLabel}生成失败`;
|
||
const description = isCompleted
|
||
? "生成结果已同步到任务历史，可继续编辑或保存到资产库。"
|
||
: String(task.error || "任务处理失败，请稍后重试。").slice(0, 500);
|
||
|
||
await pool.query(
|
||
`
|
||
INSERT INTO web_notifications (
|
||
user_id, type, title, description, target_type, target_id, metadata_json
|
||
)
|
||
SELECT $1::integer, $2::varchar, $3::varchar, $4::text, 'generation_task', $5::varchar, $6::text
|
||
WHERE NOT EXISTS (
|
||
SELECT 1
|
||
FROM web_notifications
|
||
WHERE user_id = $1
|
||
AND type = $2::varchar
|
||
AND target_type = 'generation_task'
|
||
AND target_id = $5::varchar
|
||
)
|
||
`,
|
||
[
|
||
task.user_id,
|
||
noticeType,
|
||
title,
|
||
description,
|
||
String(task.id),
|
||
JSON.stringify({ taskType: task.type, resultUrl: task.result_url || null }),
|
||
],
|
||
);
|
||
}
|
||
|
||
async function updateTaskInDb(taskId, updates) {
|
||
const nextUpdates = { ...updates };
|
||
|
||
const fields = [];
|
||
const values = [];
|
||
let idx = 1;
|
||
|
||
if (nextUpdates.status !== undefined) { fields.push(`status = $${idx++}`); values.push(nextUpdates.status); }
|
||
if (nextUpdates.progress !== undefined) {
|
||
const progress = normalizeTaskProgress(nextUpdates.progress);
|
||
if (progress !== undefined) { fields.push(`progress = $${idx++}`); values.push(progress); }
|
||
}
|
||
if (nextUpdates.resultUrl !== undefined) { fields.push(`result_url = $${idx++}`); values.push(nextUpdates.resultUrl); }
|
||
if (nextUpdates.error !== undefined) { fields.push(`error = $${idx++}`); values.push(nextUpdates.error); }
|
||
if (nextUpdates.providerTaskId !== undefined) { fields.push(`provider_task_id = $${idx++}`); values.push(nextUpdates.providerTaskId); }
|
||
if (nextUpdates.status === "completed" || nextUpdates.status === "failed") {
|
||
fields.push("completed_at = NOW()");
|
||
}
|
||
fields.push("updated_at = NOW()");
|
||
|
||
if (fields.length === 0) return;
|
||
values.push(taskId);
|
||
const { rows } = await pool.query(
|
||
`UPDATE generation_tasks SET ${fields.join(", ")} WHERE id = $${idx} RETURNING *`,
|
||
values,
|
||
);
|
||
let updatedTask = rows[0];
|
||
|
||
if (updatedTask) {
|
||
await publishTaskEvent(formatTaskEvent(updatedTask));
|
||
}
|
||
|
||
if (nextUpdates.status === "completed" && updatedTask?.result_url) {
|
||
persistTaskResultUrlToOssInBackground(updatedTask);
|
||
}
|
||
|
||
if (nextUpdates.status === "completed" || nextUpdates.status === "failed") {
|
||
await createTaskLifecycleNotification(updatedTask).catch((err) => {
|
||
console.error(`[aiTaskWorker] notification error for task ${taskId}:`, err.message);
|
||
});
|
||
}
|
||
|
||
if (nextUpdates.status === "failed") {
|
||
await refundTaskBillingOnFailure(taskId).catch((err) => {
|
||
console.error(`[aiTaskWorker] refund error for task ${taskId}:`, err.message);
|
||
});
|
||
}
|
||
}
|
||
|
||
function persistTaskResultUrlToOssInBackground(task) {
|
||
if (!task?.id || !task?.result_url) return;
|
||
|
||
Promise.resolve()
|
||
.then(async () => {
|
||
const durableUrl = await persistResultUrlToOss(task);
|
||
if (!durableUrl || durableUrl === task.result_url) return;
|
||
|
||
await pool.query(
|
||
"UPDATE generation_tasks SET result_url = $1, updated_at = NOW() WHERE id = $2 AND result_url = $3",
|
||
[durableUrl, task.id, task.result_url],
|
||
);
|
||
console.info(`[aiTaskWorker] task ${task.id} result persisted to OSS after completion`);
|
||
})
|
||
.catch((error) => {
|
||
console.warn(`[aiTaskWorker] background result persistence failed for task ${task.id}:`, error.message);
|
||
});
|
||
}
|
||
|
||
function asObject(value) {
|
||
return value && typeof value === "object" && !Array.isArray(value) ? value : undefined;
|
||
}
|
||
|
||
function asString(value) {
|
||
return typeof value === "string" && value.trim() ? value.trim() : undefined;
|
||
}
|
||
|
||
function firstString(...values) {
|
||
for (const value of values) {
|
||
const stringValue = asString(value);
|
||
if (stringValue) return stringValue;
|
||
}
|
||
return undefined;
|
||
}
|
||
|
||
function mediaExtensionFromContentType(contentType, fallbackUrl, taskType) {
|
||
const mime = String(contentType || "").split(";")[0].trim().toLowerCase();
|
||
if (mime === "image/jpeg") return "jpg";
|
||
if (mime === "image/png") return "png";
|
||
if (mime === "image/webp") return "webp";
|
||
if (mime === "image/gif") return "gif";
|
||
if (mime === "video/webm") return "webm";
|
||
if (mime === "video/quicktime") return "mov";
|
||
if (mime === "video/mp4") return "mp4";
|
||
|
||
try {
|
||
const matched = new URL(fallbackUrl).pathname.match(/\.([a-z0-9]{2,5})$/i);
|
||
if (matched?.[1]) return matched[1].toLowerCase();
|
||
} catch {
|
||
// Keep the type fallback below.
|
||
}
|
||
|
||
return taskType === "video" ? "mp4" : "png";
|
||
}
|
||
|
||
function isErrorDocumentContentType(contentType) {
|
||
return /(?:application|text)\/(?:json|xml|html|plain)|\+xml/i.test(String(contentType || ""));
|
||
}
|
||
|
||
function isOwnPersistedResultUrl(value) {
|
||
return /\/users\/[^/]+\/generation-results\//i.test(String(value || ""));
|
||
}
|
||
|
||
async function persistIncomingResultUrl(taskId, resultUrl) {
|
||
const normalizedUrl = String(resultUrl || "").trim();
|
||
if (!/^https?:\/\//i.test(normalizedUrl) || isOwnPersistedResultUrl(normalizedUrl) || !isOssConfigured()) {
|
||
return null;
|
||
}
|
||
|
||
try {
|
||
const { rows } = await pool.query(
|
||
"SELECT id, user_id, type, result_url FROM generation_tasks WHERE id = $1",
|
||
[taskId],
|
||
);
|
||
const task = rows[0];
|
||
if (!task?.user_id || !task?.type) return null;
|
||
return persistResultUrlToOss({ ...task, result_url: normalizedUrl });
|
||
} catch (error) {
|
||
console.warn(`[aiTaskWorker] result pre-persistence skipped for task ${taskId}:`, error.message);
|
||
return null;
|
||
}
|
||
}
|
||
|
||
async function persistResultUrlToOss(task) {
|
||
const resultUrl = String(task?.result_url || "").trim();
|
||
if (!/^https?:\/\//i.test(resultUrl) || isOwnPersistedResultUrl(resultUrl) || !isOssConfigured()) {
|
||
return null;
|
||
}
|
||
|
||
try {
|
||
const response = await fetch(resultUrl, { method: "GET" });
|
||
if (!response.ok) {
|
||
throw new Error(`result fetch returned ${response.status}`);
|
||
}
|
||
|
||
const contentType = response.headers.get("content-type") || (task.type === "video" ? "video/mp4" : "image/png");
|
||
if (isErrorDocumentContentType(contentType)) {
|
||
const text = await response.text().catch(() => "");
|
||
throw new Error(`result fetch returned error document: ${text.slice(0, 120)}`);
|
||
}
|
||
|
||
const buffer = Buffer.from(await response.arrayBuffer());
|
||
if (!buffer.length) {
|
||
throw new Error("result fetch returned empty content");
|
||
}
|
||
|
||
const safeUserId = String(task.user_id).replace(/[^a-zA-Z0-9_-]/g, "");
|
||
const extension = mediaExtensionFromContentType(contentType, resultUrl, task.type);
|
||
const objectKey = `users/${safeUserId}/generation-results/${task.id}-${Date.now()}-${crypto.randomUUID()}.${extension}`;
|
||
const uploaded = await putObject(objectKey, buffer, contentType, { "x-oss-object-acl": "public-read" });
|
||
return uploaded.url;
|
||
} catch (error) {
|
||
console.warn(`[aiTaskWorker] result persistence skipped for task ${task?.id}:`, error.message);
|
||
return null;
|
||
}
|
||
}
|
||
|
||
function normalizeImageResultValue(value) {
|
||
const stringValue = firstString(value);
|
||
if (!stringValue) return undefined;
|
||
if (/^(https?:)?\/\//i.test(stringValue) || /^data:image\//i.test(stringValue)) return stringValue;
|
||
const markdownImageMatch = stringValue.match(/!\[[^\]]*]\((https?:\/\/[^)\s]+)\)/i);
|
||
if (markdownImageMatch?.[1]) return markdownImageMatch[1];
|
||
const urlMatch = stringValue.match(/https?:\/\/[^\s"'<>）)]+/i);
|
||
if (urlMatch?.[0]) return urlMatch[0];
|
||
try {
|
||
const parsed = JSON.parse(stringValue);
|
||
const parsedResult = firstImageResult(parsed);
|
||
if (parsedResult) return parsedResult;
|
||
} catch {}
|
||
if (/^[A-Za-z0-9+/]+={0,2}$/.test(stringValue) && stringValue.length > 128) {
|
||
return `data:image/png;base64,${stringValue}`;
|
||
}
|
||
return undefined;
|
||
}
|
||
|
||
function firstImageResult(...values) {
|
||
for (const value of values) {
|
||
if (Array.isArray(value)) {
|
||
for (const item of value) {
|
||
const result = firstImageResult(item);
|
||
if (result) return result;
|
||
}
|
||
continue;
|
||
}
|
||
if (value && typeof value === "object") {
|
||
const result = firstImageResult(
|
||
value.url,
|
||
value.image_url,
|
||
value.imageUrl,
|
||
value.result_url,
|
||
value.resultUrl,
|
||
value.output_url,
|
||
value.outputUrl,
|
||
value.b64_image,
|
||
value.b64_json,
|
||
value.base64,
|
||
value.image,
|
||
value.content,
|
||
value.text,
|
||
value.message,
|
||
);
|
||
if (result) return result;
|
||
continue;
|
||
}
|
||
const result = normalizeImageResultValue(value);
|
||
if (result) return result;
|
||
}
|
||
return undefined;
|
||
}
|
||
|
||
function pickDeep(value, keys) {
|
||
if (!value || typeof value !== "object") return undefined;
|
||
const obj = value;
|
||
for (const key of keys) {
|
||
const direct = firstString(obj[key]);
|
||
if (direct) return direct;
|
||
}
|
||
for (const child of Object.values(obj)) {
|
||
if (child && typeof child === "object") {
|
||
const nested = pickDeep(child, keys);
|
||
if (nested) return nested;
|
||
}
|
||
}
|
||
return undefined;
|
||
}
|
||
|
||
function normalizeStatus(value) {
|
||
return String(value || "").trim().toLowerCase();
|
||
}
|
||
|
||
function isCompletedStatus(status) {
|
||
return ["completed", "complete", "succeeded", "success", "succeed", "done", "finished", "successed"].includes(status);
|
||
}
|
||
|
||
function isFailedStatus(status) {
|
||
return ["failed", "failure", "fail", "canceled", "cancelled", "expired", "error", "violation"].includes(status);
|
||
}
|
||
|
||
function extractProviderTaskId(json) {
|
||
const data = asObject(json?.data) || json;
|
||
const output = asObject(json?.output) || asObject(data?.output);
|
||
return firstString(
|
||
output?.task_id,
|
||
output?.taskId,
|
||
data?.task_id,
|
||
data?.taskId,
|
||
data?.id,
|
||
json?.task_id,
|
||
json?.taskId,
|
||
json?.id,
|
||
);
|
||
}
|
||
|
||
function extractImageUrl(json) {
|
||
const rawData = json?.data;
|
||
const data = asObject(rawData) || json;
|
||
const rawOutput = json?.output ?? data?.output;
|
||
const output = asObject(rawOutput) || asObject(data?.output);
|
||
const rawResult = data?.result ?? json?.result ?? output?.result;
|
||
const result = asObject(rawResult);
|
||
const choices =
|
||
(Array.isArray(output?.choices) && output.choices) ||
|
||
(Array.isArray(data?.choices) && data.choices) ||
|
||
(Array.isArray(json?.choices) && json.choices) ||
|
||
[];
|
||
const firstChoice = asObject(choices[0]);
|
||
const message = asObject(firstChoice?.message);
|
||
const content = Array.isArray(message?.content) ? message.content : [];
|
||
const firstContent = asObject(content[0]);
|
||
const outputResults = Array.isArray(output?.results) ? output.results : [];
|
||
const topLevelResults = Array.isArray(json?.results) ? json.results : [];
|
||
const dataResults = Array.isArray(data?.results) ? data.results : [];
|
||
const resultResults = Array.isArray(result?.results) ? result.results : [];
|
||
const dataImages = Array.isArray(data?.images) ? data.images : [];
|
||
const dataImageUrls = Array.isArray(data?.image_urls) ? data.image_urls : [];
|
||
const dataUrls = Array.isArray(data?.urls) ? data.urls : [];
|
||
const outputImages = Array.isArray(output?.images) ? output.images : [];
|
||
const outputImageUrls = Array.isArray(output?.image_urls) ? output.image_urls : [];
|
||
const resultImages = Array.isArray(result?.images) ? result.images : [];
|
||
const resultUrls = Array.isArray(result?.urls) ? result.urls : [];
|
||
const candidates = [
|
||
...topLevelResults,
|
||
...dataResults,
|
||
...outputResults,
|
||
...resultResults,
|
||
...dataImages,
|
||
...dataImageUrls,
|
||
...dataUrls,
|
||
...outputImages,
|
||
...outputImageUrls,
|
||
...resultImages,
|
||
...resultUrls,
|
||
];
|
||
|
||
return firstImageResult(
|
||
rawData,
|
||
rawOutput,
|
||
rawResult,
|
||
firstContent?.image,
|
||
firstContent?.image_url,
|
||
firstContent?.image_url?.url,
|
||
message?.content,
|
||
firstChoice?.delta?.content,
|
||
candidates,
|
||
data?.image_url,
|
||
data?.imageUrl,
|
||
data?.result_url,
|
||
data?.resultUrl,
|
||
data?.output_url,
|
||
data?.outputUrl,
|
||
output?.image_url,
|
||
output?.imageUrl,
|
||
output?.result_url,
|
||
output?.resultUrl,
|
||
output?.output_url,
|
||
output?.outputUrl,
|
||
result?.image_url,
|
||
result?.imageUrl,
|
||
result?.result_url,
|
||
result?.resultUrl,
|
||
result?.output_url,
|
||
result?.outputUrl,
|
||
pickDeep(json, ["image", "image_url", "imageUrl", "result_url", "resultUrl", "output_url", "outputUrl", "url", "b64_image", "b64_json", "base64"]),
|
||
);
|
||
}
|
||
|
||
function extractGeminiImageUrl(json) {
|
||
// Gemini response: candidates[].content.parts[].inlineData (base64) or text (URL)
|
||
const candidates = Array.isArray(json?.candidates) ? json.candidates : [];
|
||
for (const candidate of candidates) {
|
||
const parts = Array.isArray(candidate?.content?.parts) ? candidate.content.parts : [];
|
||
for (const part of parts) {
|
||
const inlineData = part?.inlineData;
|
||
if (inlineData?.data) {
|
||
const mimeType = inlineData.mimeType || "image/png";
|
||
return `data:${mimeType};base64,${inlineData.data}`;
|
||
}
|
||
}
|
||
}
|
||
// Also check for direct URL in candidate text
|
||
for (const candidate of candidates) {
|
||
const parts = Array.isArray(candidate?.content?.parts) ? candidate.content.parts : [];
|
||
for (const part of parts) {
|
||
if (part?.text && /^https?:\/\/.+\.(png|jpg|jpeg|webp|gif)/i.test(part.text)) {
|
||
return part.text;
|
||
}
|
||
}
|
||
}
|
||
return null;
|
||
}
|
||
|
||
function extractVideoUrl(json) {
|
||
const data = asObject(json?.data) || json;
|
||
const output = asObject(json?.output) || asObject(data?.output);
|
||
const result = asObject(data?.task_result) || asObject(json?.task_result) || asObject(output?.task_result);
|
||
const videos = Array.isArray(result?.videos) ? result.videos : [];
|
||
const firstVideo = asObject(videos[0]);
|
||
|
||
return firstString(
|
||
output?.video_url,
|
||
output?.output_video_url,
|
||
output?.outputVideoUrl,
|
||
output?.watermark_video_url,
|
||
data?.video_url,
|
||
data?.videoUrl,
|
||
data?.output_video_url,
|
||
data?.outputVideoUrl,
|
||
data?.url,
|
||
data?.result?.url,
|
||
firstVideo?.url,
|
||
firstVideo?.video_url,
|
||
pickDeep(json, ["video_url", "output_video_url", "outputVideoUrl", "watermark_video_url", "videoUrl", "download_url", "downloadUrl", "content_url", "contentUrl", "url"]),
|
||
);
|
||
}
|
||
|
||
function extractErrorMessage(json, fallback) {
|
||
const data = asObject(json?.data) || {};
|
||
const output = asObject(json?.output) || {};
|
||
const error = asObject(json?.error) || {};
|
||
return firstString(
|
||
output.message,
|
||
output.code,
|
||
data.task_status_msg,
|
||
data.failure_reason,
|
||
data.message,
|
||
error.message,
|
||
error.error,
|
||
json?.message,
|
||
json?.error,
|
||
) || fallback;
|
||
}
|
||
|
||
async function fetchJson(url, headers) {
|
||
const res = await fetch(url, { method: "GET", headers });
|
||
if (!res.ok) return { ok: false, json: null };
|
||
return { ok: true, json: await res.json() };
|
||
}
|
||
|
||
async function pollGrsaiImage(_taskId, providerTaskId, apiKey, baseUrl, resultEndpoint) {
|
||
const url = `${baseUrl}${resultEndpoint}?id=${encodeURIComponent(providerTaskId)}`;
|
||
const { ok, json } = await fetchJson(url, {
|
||
Authorization: `Bearer ${apiKey}`,
|
||
Accept: "application/json",
|
||
});
|
||
if (!ok) {
|
||
console.warn(`[grsai-poll] task ${_taskId} fetch not ok, url=${url}`);
|
||
return { status: "running", progress: 50 };
|
||
}
|
||
|
||
const data = asObject(json?.data) || json;
|
||
const output = asObject(data?.output) || asObject(json?.output) || {};
|
||
const status = normalizeStatus(
|
||
output.task_status ||
|
||
output.status ||
|
||
data.task_status ||
|
||
data.status ||
|
||
data.state ||
|
||
json?.task_status ||
|
||
json?.status ||
|
||
json?.state,
|
||
);
|
||
const resultUrl = extractImageUrl(json);
|
||
console.info(`[grsai-poll] task ${_taskId} status=${status} resultUrl=${resultUrl ? "yes" : "no"} raw=${JSON.stringify(json).slice(0, 300)}`);
|
||
if (resultUrl) {
|
||
return { status: "completed", progress: 100, resultUrl };
|
||
}
|
||
if (isCompletedStatus(status)) {
|
||
const completedUrl = extractImageUrl(json);
|
||
if (!completedUrl) return { status: "failed", error: "Image generation completed without a result url" };
|
||
return { status: "completed", progress: 100, resultUrl: completedUrl };
|
||
}
|
||
if (isFailedStatus(status)) {
|
||
return { status: "failed", error: extractErrorMessage(json, "Image generation failed") };
|
||
}
|
||
return { status: "running", progress: Math.min(90, 30 + Math.random() * 40) };
|
||
}
|
||
|
||
async function pollDashscopeImage(_taskId, providerTaskId, apiKey) {
|
||
const url = `https://dashscope.aliyuncs.com/api/v1/tasks/${encodeURIComponent(providerTaskId)}`;
|
||
const { ok, json } = await fetchJson(url, {
|
||
Authorization: `Bearer ${apiKey}`,
|
||
Accept: "application/json",
|
||
});
|
||
if (!ok) return { status: "running", progress: 50 };
|
||
|
||
const output = asObject(json?.output) || {};
|
||
const status = normalizeStatus(output.task_status || json?.task_status);
|
||
const resultUrl = extractImageUrl(json);
|
||
if (isCompletedStatus(status)) {
|
||
if (!resultUrl) return { status: "failed", error: "DashScope image generation completed without a result url" };
|
||
return { status: "completed", progress: 100, resultUrl };
|
||
}
|
||
if (isFailedStatus(status)) {
|
||
return { status: "failed", error: extractErrorMessage(json, "DashScope image generation failed") };
|
||
}
|
||
return { status: "running", progress: Math.min(90, 30 + Math.random() * 40) };
|
||
}
|
||
|
||
function base64Url(input) {
|
||
return Buffer.from(input)
|
||
.toString("base64")
|
||
.replace(/\+/g, "-")
|
||
.replace(/\//g, "_")
|
||
.replace(/=+$/, "");
|
||
}
|
||
|
||
function parseKlingCredential(apiKey) {
|
||
const raw = String(apiKey || "");
|
||
const colonIndex = raw.indexOf(":");
|
||
if (colonIndex <= 0) return null;
|
||
const accessKey = raw.slice(0, colonIndex).trim();
|
||
const secretKey = raw.slice(colonIndex + 1).trim();
|
||
return accessKey && secretKey ? { accessKey, secretKey } : null;
|
||
}
|
||
|
||
function createKlingJwt(accessKey, secretKey) {
|
||
const header = { alg: "HS256", typ: "JWT" };
|
||
const now = Math.floor(Date.now() / 1000);
|
||
const payload = { iss: accessKey, exp: now + 1800, nbf: now - 5 };
|
||
const unsigned = `${base64Url(JSON.stringify(header))}.${base64Url(JSON.stringify(payload))}`;
|
||
const signature = crypto.createHmac("sha256", secretKey).update(unsigned).digest("base64")
|
||
.replace(/\+/g, "-")
|
||
.replace(/\//g, "_")
|
||
.replace(/=+$/, "");
|
||
return `${unsigned}.${signature}`;
|
||
}
|
||
|
||
function getPollRequest(providerTaskId, apiKey, providerConfig) {
|
||
const protocol = providerConfig.protocol;
|
||
const baseUrl = providerConfig.baseUrl || "";
|
||
|
||
if (
|
||
protocol === "wan-i2v" ||
|
||
protocol === "wan-s2v" ||
|
||
protocol === "kling-dashscope" ||
|
||
String(protocol || "").startsWith("happyhorse-")
|
||
) {
|
||
return {
|
||
url: `${baseUrl}/api/v1/tasks/${encodeURIComponent(providerTaskId)}`,
|
||
headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
|
||
};
|
||
}
|
||
|
||
if (protocol === "kling-omni") {
|
||
const credential = parseKlingCredential(apiKey);
|
||
return {
|
||
url: `${baseUrl}/v1/videos/omni-video/${encodeURIComponent(providerTaskId)}`,
|
||
headers: {
|
||
Authorization: `Bearer ${credential ? createKlingJwt(credential.accessKey, credential.secretKey) : apiKey}`,
|
||
Accept: "application/json",
|
||
},
|
||
};
|
||
}
|
||
|
||
return {
|
||
url: `${baseUrl}${providerConfig.endpoint}/${encodeURIComponent(providerTaskId)}`,
|
||
headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
|
||
};
|
||
}
|
||
|
||
async function pollVideoTask(_taskId, providerTaskId, apiKey, providerConfig) {
|
||
const { url, headers } = getPollRequest(providerTaskId, apiKey, providerConfig);
|
||
const { ok, json } = await fetchJson(url, headers);
|
||
if (!ok) return { status: "running", progress: 50 };
|
||
|
||
const data = asObject(json?.data) || json;
|
||
const output = asObject(json?.output) || {};
|
||
const status = normalizeStatus(
|
||
output.task_status ||
|
||
data.task_status ||
|
||
data.status ||
|
||
json?.task_status ||
|
||
json?.status,
|
||
);
|
||
|
||
const resultUrl = extractVideoUrl(json);
|
||
if (isCompletedStatus(status) || resultUrl) {
|
||
return { status: "completed", progress: 100, resultUrl: resultUrl || null };
|
||
}
|
||
if (isFailedStatus(status)) {
|
||
return { status: "failed", error: extractErrorMessage(json, "Video generation failed") };
|
||
}
|
||
const progress = Number(data.progress || output.progress);
|
||
return { status: "running", progress: Number.isFinite(progress) ? Math.min(95, progress) : Math.min(90, 30 + Math.random() * 30) };
|
||
}
|
||
|
||
function getMaxPollAttempts(type, providerConfig) {
|
||
if (type === "image" && providerConfig?.transport === "grsai-image") {
|
||
return Number.isFinite(GRS_IMAGE_MAX_POLL_ATTEMPTS) && GRS_IMAGE_MAX_POLL_ATTEMPTS > 0
|
||
? Math.trunc(GRS_IMAGE_MAX_POLL_ATTEMPTS)
|
||
: 40;
|
||
}
|
||
if (type === "video") return 400;
|
||
return MAX_POLL_ATTEMPTS;
|
||
}
|
||
|
||
function startPolling(taskDbId, { providerTaskId, apiKey, type, providerConfig, leaseToken, keyManager, onTaskFailed, skipPersist = false }) {
|
||
if (activePollers.has(taskDbId)) return;
|
||
if (!skipPersist) {
|
||
persistPollerState(taskDbId, { providerTaskId, type, providerConfig, leaseToken }).catch((err) => {
|
||
console.error(`[aiTaskWorker] failed to persist poller state for task ${taskDbId}:`, err.message);
|
||
});
|
||
}
|
||
|
||
let attempts = 0;
|
||
const maxPollAttempts = getMaxPollAttempts(type, providerConfig);
|
||
const interval = setInterval(async () => {
|
||
attempts++;
|
||
if (attempts > maxPollAttempts) {
|
||
clearInterval(interval);
|
||
activePollers.delete(taskDbId);
|
||
if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
|
||
if (typeof onTaskFailed === "function") {
|
||
const handled = await onTaskFailed("Task timed out").catch((fallbackErr) => {
|
||
console.error(`[aiTaskWorker] fallback error for task ${taskDbId}:`, fallbackErr.message);
|
||
return false;
|
||
});
|
||
if (handled) return;
|
||
}
|
||
await updateTaskInDb(taskDbId, { status: "failed", error: "Task timed out" });
|
||
await clearPollerState(taskDbId).catch(() => {});
|
||
return;
|
||
}
|
||
|
||
try {
|
||
// Check if task was cancelled by user
|
||
const { rows: [taskRow] } = await pool.query("SELECT status FROM generation_tasks WHERE id = $1", [taskDbId]);
|
||
if (!taskRow || taskRow.status === "cancelled") {
|
||
clearInterval(interval);
|
||
activePollers.delete(taskDbId);
|
||
await clearPollerState(taskDbId).catch(() => {});
|
||
if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
|
||
return;
|
||
}
|
||
await refreshPollerHeartbeat(taskDbId).catch(() => {});
|
||
|
||
let result;
|
||
if (type === "image") {
|
||
if (providerConfig.transport === "dashscope-image") {
|
||
result = await pollDashscopeImage(taskDbId, providerTaskId, apiKey);
|
||
} else {
|
||
result = await pollGrsaiImage(taskDbId, providerTaskId, apiKey, providerConfig.baseUrl, providerConfig.resultEndpoint || "/result");
|
||
}
|
||
} else {
|
||
result = await pollVideoTask(taskDbId, providerTaskId, apiKey, providerConfig);
|
||
}
|
||
|
||
if (result.status === "completed" || result.status === "failed") {
|
||
clearInterval(interval);
|
||
activePollers.delete(taskDbId);
|
||
if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
|
||
if (result.status === "failed" && typeof onTaskFailed === "function") {
|
||
const handled = await onTaskFailed(result.error || "Task failed").catch((fallbackErr) => {
|
||
console.error(`[aiTaskWorker] fallback error for task ${taskDbId}:`, fallbackErr.message);
|
||
return false;
|
||
});
|
||
if (handled) return;
|
||
}
|
||
}
|
||
|
||
await updateTaskInDb(taskDbId, result);
|
||
if (result.status === "completed" || result.status === "failed") {
|
||
await clearPollerState(taskDbId).catch(() => {});
|
||
}
|
||
} catch (err) {
|
||
console.error(`[aiTaskWorker] poll error for task ${taskDbId}:`, err.message);
|
||
}
|
||
}, POLL_INTERVAL_MS);
|
||
|
||
activePollers.set(taskDbId, { interval, leaseToken });
|
||
}
|
||
|
||
function stopPolling(taskDbId) {
|
||
const poller = activePollers.get(taskDbId);
|
||
if (poller) {
|
||
clearInterval(poller.interval);
|
||
activePollers.delete(taskDbId);
|
||
}
|
||
clearPollerState(taskDbId).catch(() => {});
|
||
}
|
||
|
||
function getActiveCount() {
|
||
return activePollers.size;
|
||
}
|
||
|
||
async function recoverRunnablePollers() {
|
||
await ensureTaskPollerStore();
|
||
const staleInterval = `${Math.max(5, Math.ceil(POLLER_OWNER_STALE_MS / 1000))} seconds`;
|
||
const { rows } = await pool.query(
|
||
`
|
||
SELECT p.task_id
|
||
FROM generation_task_pollers p
|
||
JOIN generation_tasks t ON t.id = p.task_id
|
||
WHERE t.status IN ('pending', 'running')
|
||
AND (
|
||
p.owner_heartbeat_at IS NULL
|
||
OR p.owner_heartbeat_at < NOW() - ($1::text)::interval
|
||
)
|
||
ORDER BY p.owner_heartbeat_at NULLS FIRST, p.updated_at ASC
|
||
LIMIT 20
|
||
`,
|
||
[staleInterval],
|
||
);
|
||
|
||
for (const row of rows) {
|
||
const taskId = row.task_id;
|
||
if (activePollers.has(taskId)) continue;
|
||
const poller = await claimPoller(taskId);
|
||
if (!poller || activePollers.has(taskId)) continue;
|
||
|
||
const apiKey = await getLeaseKey(poller.lease_token);
|
||
if (apiKey == null) {
|
||
console.warn(`[aiTaskWorker] cannot recover task ${taskId}: active lease not found`);
|
||
continue;
|
||
}
|
||
|
||
console.info(`[aiTaskWorker] recovering poller for task ${taskId}`);
|
||
startPolling(taskId, {
|
||
providerTaskId: poller.provider_task_id,
|
||
apiKey,
|
||
type: poller.task_type,
|
||
providerConfig: parseProviderConfig(poller.provider_config_json),
|
||
leaseToken: poller.lease_token,
|
||
keyManager: require("./keyManager"),
|
||
skipPersist: true,
|
||
});
|
||
}
|
||
}
|
||
|
||
// --- Periodic stale task cleanup ---
|
||
// Runs every 5 minutes, marks tasks stuck in 'pending'/'running' for too long as 'failed'.
|
||
// This catches cases where the worker crashed, the provider API never responded,
|
||
// or the cancel request failed silently on the client side.
|
||
const STALE_TASK_CLEANUP_INTERVAL_MS = 5 * 60 * 1000;
|
||
let staleTaskCleanupTimer = null;
|
||
|
||
async function runStaleTaskCleanup() {
|
||
try {
|
||
const { rows } = await pool.query(
|
||
`UPDATE generation_tasks
|
||
SET status = 'failed', error = '任务超时自动释放', updated_at = NOW()
|
||
WHERE status IN ('pending', 'running')
|
||
AND GREATEST(updated_at, COALESCE(last_poll_at, created_at)) < NOW() - INTERVAL '10 minutes'
|
||
RETURNING id`,
|
||
);
|
||
for (const row of rows) {
|
||
await publishTaskEvent({
|
||
taskId: row.id,
|
||
status: "failed",
|
||
progress: null,
|
||
resultUrl: null,
|
||
error: "任务超时自动释放",
|
||
});
|
||
// Also stop any active poller for this task
|
||
const poller = activePollers.get(row.id);
|
||
if (poller) {
|
||
clearInterval(poller.interval);
|
||
activePollers.delete(row.id);
|
||
}
|
||
await clearPollerState(row.id).catch(() => {});
|
||
}
|
||
if (rows.length > 0) {
|
||
console.log(`[aiTaskWorker] Cleaned up ${rows.length} stale task(s)`);
|
||
}
|
||
} catch (err) {
|
||
console.error("[aiTaskWorker] Stale task cleanup failed:", err.message);
|
||
}
|
||
}
|
||
|
||
async function startTaskEventListener() {
|
||
if (taskEventListenerClient) return;
|
||
if (taskEventListenerStarting) return taskEventListenerStarting;
|
||
|
||
taskEventListenerStarting = (async () => {
|
||
const client = await pool.connect();
|
||
let released = false;
|
||
|
||
const releaseClient = () => {
|
||
if (released) return;
|
||
released = true;
|
||
taskEventListenerClient = null;
|
||
try {
|
||
client.release();
|
||
} catch {}
|
||
};
|
||
|
||
client.on("notification", (message) => {
|
||
if (message.channel !== TASK_EVENT_CHANNEL || !message.payload) return;
|
||
try {
|
||
const payload = JSON.parse(message.payload);
|
||
if (payload?.origin === TASK_EVENT_ORIGIN) return;
|
||
emitTaskEvent(payload?.event || payload);
|
||
} catch (err) {
|
||
console.error("[aiTaskWorker] task event notification parse failed:", err.message);
|
||
}
|
||
});
|
||
|
||
client.on("error", (err) => {
|
||
console.error("[aiTaskWorker] task event listener error:", err.message);
|
||
releaseClient();
|
||
setTimeout(() => {
|
||
startTaskEventListener().catch((restartErr) => {
|
||
console.error("[aiTaskWorker] task event listener restart failed:", restartErr.message);
|
||
});
|
||
}, 5000).unref?.();
|
||
});
|
||
|
||
await client.query(`LISTEN ${TASK_EVENT_CHANNEL}`);
|
||
taskEventListenerClient = client;
|
||
console.log(`[aiTaskWorker] listening for task events on ${TASK_EVENT_CHANNEL}`);
|
||
})();
|
||
|
||
try {
|
||
await taskEventListenerStarting;
|
||
} finally {
|
||
taskEventListenerStarting = null;
|
||
}
|
||
}
|
||
|
||
async function stopTaskEventListener() {
|
||
const client = taskEventListenerClient;
|
||
taskEventListenerClient = null;
|
||
if (!client) return;
|
||
try {
|
||
await client.query(`UNLISTEN ${TASK_EVENT_CHANNEL}`);
|
||
} catch {}
|
||
client.release();
|
||
}
|
||
|
||
function startStaleTaskCleanup() {
|
||
if (staleTaskCleanupTimer) return;
|
||
staleTaskCleanupTimer = setInterval(runStaleTaskCleanup, STALE_TASK_CLEANUP_INTERVAL_MS);
|
||
// Run once shortly after startup
|
||
setTimeout(runStaleTaskCleanup, 10_000);
|
||
}
|
||
|
||
function stopStaleTaskCleanup() {
|
||
if (staleTaskCleanupTimer) {
|
||
clearInterval(staleTaskCleanupTimer);
|
||
staleTaskCleanupTimer = null;
|
||
}
|
||
}
|
||
|
||
function startPollerRecovery() {
|
||
if (pollerRecoveryTimer) return;
|
||
ensureTaskPollerStore()
|
||
.then(() => recoverRunnablePollers())
|
||
.catch((err) => console.error("[aiTaskWorker] initial poller recovery failed:", err.message));
|
||
pollerRecoveryTimer = setInterval(() => {
|
||
recoverRunnablePollers().catch((err) => {
|
||
console.error("[aiTaskWorker] poller recovery failed:", err.message);
|
||
});
|
||
}, POLLER_RECOVERY_INTERVAL_MS);
|
||
}
|
||
|
||
function stopPollerRecovery() {
|
||
if (pollerRecoveryTimer) {
|
||
clearInterval(pollerRecoveryTimer);
|
||
pollerRecoveryTimer = null;
|
||
}
|
||
}
|
||
|
||
module.exports = {
|
||
startPolling,
|
||
stopPolling,
|
||
updateTaskInDb,
|
||
getActiveCount,
|
||
extractProviderTaskId,
|
||
extractImageUrl,
|
||
extractGeminiImageUrl,
|
||
extractVideoUrl,
|
||
parseKlingCredential,
|
||
createKlingJwt,
|
||
taskEvents,
|
||
startTaskEventListener,
|
||
stopTaskEventListener,
|
||
startPollerRecovery,
|
||
stopPollerRecovery,
|
||
startStaleTaskCleanup,
|
||
stopStaleTaskCleanup,
|
||
};
|