Files
atlas-seller-kit/apps/worker/src/media-ai-edit.ts
T

457 lines
24 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import { createHash, randomUUID } from 'node:crypto';
import { lookup } from 'node:dns/promises';
import { request as httpsRequest } from 'node:https';
import { isIP, type LookupFunction } from 'node:net';
import { GetObjectCommand, PutObjectCommand, type S3Client } from '@aws-sdk/client-s3';
import type { MediaAiEditFailureStage, MediaAiEditSettings } from '@atlas/contracts';
import type { Database } from '@atlas/database';
import sharp from 'sharp';
import { buildMediaAiPrompt } from './media-ai-prompt.js';
import type { OutboxJob } from './outbox-dispatcher.js';
type StoredAsset = { id: string; storage_provider: string; bucket: string; object_key: string };
type EditTask = {
status: 'queued' | 'running' | 'succeeded' | 'failed' | 'cancelled';
source_asset_id: string;
reference_asset_ids: string[];
settings: MediaAiEditSettings;
provider_task_id: string | null;
created_at: Date | string;
};
const sizes: Record<MediaAiEditSettings['ratio'], [number, number]> = {
'3:4': [1152, 1536], '1:1': [1024, 1024], '4:5': [1280, 1600],
};
type AiEditFailureCode = 'AI_IMAGE_GENERATION_FAILED' | 'AI_IMAGE_CONTENT_REJECTED' | 'AI_IMAGE_PROVIDER_FAILED';
export type AiEditFailure = {
code: AiEditFailureCode;
source: 'pcore' | 'atlas';
stage: MediaAiEditFailureStage;
detail: string;
providerErrorRaw: string | null;
};
class AiEditStageError extends Error {
public constructor(public readonly stage: MediaAiEditFailureStage, public readonly cause: unknown) {
super(cause instanceof Error ? cause.message : '未知错误');
}
}
async function atStage<T>(stage: MediaAiEditFailureStage, action: () => Promise<T>): Promise<T> {
try { return await action(); }
catch (error) { throw error instanceof AiEditStageError ? error : new AiEditStageError(stage, error); }
}
/** 只保存可供排查的短摘要,避免把密钥或带签名的地址写入任务结果。 */
function safeFailureDetail(error: unknown, stage: MediaAiEditFailureStage): string {
const message = error instanceof Error ? error.message : '未知错误';
if (stage === 'task_setup' || stage === 'input_preparation' || stage === 'result_storage') {
if (['图片对象不存在', '源图片超过处理上限', '参考图片格式不受支持',
'参考图片素材不存在'].includes(message)) return message;
if (error instanceof PcoreResponseError) return message.slice(0, 500);
const code = error && typeof error === 'object' && 'code' in error && typeof error.code === 'string'
&& /^[A-Z0-9_]{2,32}$/i.test(error.code) ? `(${error.code})` : '';
return `${error instanceof Error ? error.name : 'Error'}${code}`;
}
return message.replace(/Bearer\s+\S+/gi, 'Bearer [已隐藏]')
.replace(/https?:\/\/[^\s)]+/gi, '[地址已隐藏]')
.replace(/(?:api[_-]?key|token|secret)\s*[:=]\s*\S+/gi, '[凭据已隐藏]')
.replace(/\bai-edits\/[^\s)]+/gi, '[对象键已隐藏]')
.slice(0, 500);
}
export function classifyAiEditFailure(error: unknown): AiEditFailure {
const stage = error instanceof AiEditStageError ? error.stage : 'task_setup';
const cause = error instanceof AiEditStageError ? error.cause : error;
const providerError = cause instanceof PcoreResponseError ? cause : null;
const transportFailure = (stage === 'provider_submission' || stage === 'provider_polling')
&& cause instanceof Error && ['TypeError', 'TimeoutError', 'AbortError'].includes(cause.name);
return {
code: 'AI_IMAGE_GENERATION_FAILED',
source: providerError || transportFailure ? 'pcore' : 'atlas',
stage,
detail: safeFailureDetail(cause, stage),
providerErrorRaw: providerError?.rawResponse ?? null,
};
}
export class PcoreResponseError extends Error {
public readonly rawResponse: string | null;
public constructor(message: string, rawResponse: string | null) {
super(message);
this.rawResponse = rawResponse ? rawResponse.slice(0, 32_768) : null;
}
}
class PcoreTerminalError extends PcoreResponseError {
public constructor(public readonly errorCode: AiEditFailureCode, rawResponse: string | null) {
super(errorCode === 'AI_IMAGE_CONTENT_REJECTED'
? 'Pcore 未能完成这次生图,可能是输入或生成结果被审核拦截'
: 'Pcore 图片任务失败', rawResponse);
}
}
function pcoreErrorText(value: unknown): string {
if (typeof value === 'string') return value;
if (!value || typeof value !== 'object') return '';
const error = value as { code?: unknown; message?: unknown };
return [error.code, error.message].filter((item) => typeof item === 'string').join(' ');
}
function pcoreFailureCode(value: unknown): AiEditFailureCode {
if (!value || typeof value !== 'object') return 'AI_IMAGE_PROVIDER_FAILED';
const error = value as { code?: unknown; type?: unknown };
const explicitCode = [error.code, error.type].filter((item) => typeof item === 'string').join(' ');
return /safety|content[_-]?policy|moderation|blocked/i.test(explicitCode)
? 'AI_IMAGE_CONTENT_REJECTED' : 'AI_IMAGE_PROVIDER_FAILED';
}
function publicIpv4(address: string): boolean {
if (isIP(address) !== 4) return false;
const [a, b, c] = address.split('.').map(Number);
return !(a === 0 || a === 10 || a === 127 || a >= 224
|| a === 100 && b >= 64 && b <= 127
|| a === 169 && b === 254
|| a === 172 && b >= 16 && b <= 31
|| a === 192 && (b === 168 || b === 0 && (c === 0 || c === 2))
|| a === 198 && (b === 18 || b === 19 || b === 51 && c === 100)
|| a === 203 && b === 0 && c === 113);
}
async function readObject(storage: S3Client, asset: StoredAsset): Promise<Buffer> {
const response = await storage.send(new GetObjectCommand({ Bucket: asset.bucket, Key: asset.object_key }));
if (!response.Body) throw new Error('图片对象不存在');
const chunks: Buffer[] = [];
let bytes = 0;
for await (const chunk of response.Body as AsyncIterable<Uint8Array>) {
bytes += chunk.byteLength;
if (bytes > 20 * 1024 * 1024) throw new Error('源图片超过处理上限');
chunks.push(Buffer.from(chunk));
}
return Buffer.concat(chunks, bytes);
}
/** 普通编辑向供应商传原始图片字节,避免转码改变参考图内容。 */
async function referenceImageDataUri(bytes: Buffer): Promise<string> {
const image = sharp(bytes, { limitInputPixels: 100_000_000, failOn: 'error' });
const meta = await image.metadata();
if (!['jpeg', 'png', 'webp'].includes(meta.format ?? '') || !meta.width || !meta.height)
throw new Error('参考图片格式不受支持');
const mime = meta.format === 'jpeg' ? 'image/jpeg' : `image/${meta.format}`;
return `data:${mime};base64,${bytes.toString('base64')}`;
}
async function outpaintInputs(bytes: Buffer): Promise<{ image: Buffer; mask: Buffer; original: Buffer; left: number; top: number }> {
const [width, height] = sizes['3:4'];
const original = await sharp(bytes, { limitInputPixels: 100_000_000, failOn: 'error' })
.rotate().resize(width, height, { fit: 'inside', withoutEnlargement: false })
.ensureAlpha().png().toBuffer();
const meta = await sharp(original).metadata();
const left = Math.floor((width - meta.width!) / 2);
const top = Math.floor((height - meta.height!) / 2);
const canvas = await sharp({ create: { width, height, channels: 4, background: '#00000000' } })
.composite([{ input: original, left, top }]).png().toBuffer();
const opaque = await sharp({ create: { width: meta.width!, height: meta.height!, channels: 4, background: '#ffffffff' } })
.png().toBuffer();
const mask = await sharp({ create: { width, height, channels: 4, background: '#00000000' } })
.composite([{ input: opaque, left, top }]).png().toBuffer();
return {
image: canvas,
mask,
original, left, top,
};
}
function providerConfig(): { base: string; key: string } {
const key = process.env.PCORE_API_KEY;
const rawBase = process.env.PCORE_BASE_URL?.trim();
if (!key || !rawBase) throw new Error('PCORE_API_KEY 和 PCORE_BASE_URL 必须配置');
const url = new URL(rawBase);
const path = url.pathname.replace(/\/+$/, '');
if (url.protocol !== 'https:' || url.username || url.password || url.port
|| !['', '/v1'].includes(path) || url.search || url.hash)
throw new Error('PCORE_BASE_URL 必须是 HTTPS 根地址或 /v1 地址');
return { base: `${url.origin}/v1`, key };
}
async function providerRequest(url: string, key: string, body?: Record<string, unknown>): Promise<Record<string, unknown>> {
const response = await fetch(url, {
method: body ? 'POST' : 'GET',
headers: { Authorization: `Bearer ${key}`, ...(body ? { 'Content-Type': 'application/json' } : {}) },
body: body ? JSON.stringify(body) : undefined,
signal: AbortSignal.timeout(body ? 90_000 : 35_000),
redirect: 'error',
});
if (!response.ok) throw new PcoreResponseError(
`Pcore 图片任务请求失败(HTTP ${response.status})`, await response.text());
let result: Record<string, unknown>;
try {
result = await response.json() as Record<string, unknown>;
if (!result || typeof result !== 'object' || Array.isArray(result)) throw new Error('invalid response');
} catch {
throw new PcoreResponseError('Pcore 返回的图片任务响应不是有效 JSON 对象', null);
}
if (result.error && (body || !['failed', 'error', 'cancelled'].includes(String(result.status ?? ''))))
throw new PcoreResponseError(`Pcore 图片任务请求被拒绝${pcoreErrorText(result.error)
? `(${pcoreErrorText(result.error).slice(0, 200)})` : ''}`,
JSON.stringify(result.error));
return result;
}
async function submitProvider(storage: S3Client, assets: StoredAsset[], settings: MediaAiEditSettings): Promise<string> {
const provider = await atStage('provider_submission', async () => providerConfig());
const size = sizes[settings.mode === 'outpaint' ? '3:4' : settings.ratio];
const main = await atStage('input_preparation', () => readObject(storage, assets[0]));
const outpaint = settings.mode === 'outpaint'
? await atStage('input_preparation', () => outpaintInputs(main)) : null;
const images = await atStage('input_preparation', async () => outpaint
? [`data:image/png;base64,${outpaint.image.toString('base64')}`]
: await Promise.all(assets.map(async (asset, index) =>
referenceImageDataUri(index === 0 ? main : await readObject(storage, asset)))));
const body: Record<string, unknown> = {
model: 'gpt-image-2', prompt: buildMediaAiPrompt(settings), n: 1,
size: `${size[0]}x${size[1]}`, async: true,
images: images.map((image_url) => ({ image_url })),
};
if (outpaint) body.mask = { image_url: `data:image/png;base64,${outpaint.mask.toString('base64')}` };
const result = await atStage('provider_submission', () =>
providerRequest(`${provider.base}/images/edits`, provider.key, body));
const id = result.task_id ?? result.id;
if (typeof id !== 'string' || !id || id.length > 256)
throw new AiEditStageError('provider_submission', new PcoreResponseError('Pcore 未返回有效的任务 ID', null));
return id;
}
export async function downloadProviderImage(urlValue: string): Promise<Buffer> {
const url = new URL(urlValue);
if (url.protocol !== 'https:' || url.username || url.password || url.port || isIP(url.hostname))
throw new Error('Pcore 结果地址不安全');
const addresses = await lookup(url.hostname, { all: true, family: 4 });
if (!addresses.length || addresses.some((item) => !publicIpv4(item.address)))
throw new Error('Pcore 结果地址不安全');
/** 固定到校验过的公网 IPv4,同时保留域名用于 TLS 证书校验。 */
const response = await new Promise<import('node:http').IncomingMessage>((resolve, reject) => {
const request = httpsRequest(url, {
lookup: pinnedLookup(addresses[0].address),
timeout: 60_000,
}, resolve);
request.on('error', reject);
request.on('timeout', () => request.destroy(new PcoreResponseError('Pcore 结果下载超时', null)));
request.end();
});
if (response.statusCode !== 200 || Number(response.headers['content-length'] ?? 0) > 25 * 1024 * 1024) {
response.resume();
throw new PcoreResponseError(`Pcore 结果下载失败(HTTP ${response.statusCode})`, null);
}
const chunks: Buffer[] = [];
let size = 0;
for await (const chunk of response) {
size += chunk.length;
if (size > 25 * 1024 * 1024) {
response.destroy();
throw new Error('Pcore 结果超过处理上限');
}
chunks.push(Buffer.from(chunk));
}
return Buffer.concat(chunks, size);
}
export function pinnedLookup(address: string): LookupFunction {
return (_hostname, options, callback) => {
const pinned = { address, family: 4 };
if (options.all) callback(null, [pinned]);
else callback(null, pinned.address, pinned.family);
};
}
export function providerImageResult(result: Record<string, unknown>): Buffer | string {
const first = Array.isArray(result.data) ? result.data[0] as Record<string, unknown> | undefined : undefined;
for (const encoded of [first?.b64_json, first?.b64_image]) {
if (typeof encoded !== 'string' || !encoded.trim()) continue;
if (encoded.length > 35 * 1024 * 1024) throw new Error('Pcore 结果超过处理上限');
const bytes = Buffer.from(encoded.replace(/^data:[^,]*,/, ''), 'base64');
if (bytes.length) return bytes;
}
if (typeof first?.url === 'string' && first.url) return first.url;
throw new Error('Pcore 任务完成但没有图片');
}
async function pollProvider(taskId: string): Promise<Buffer | null> {
const provider = await atStage('provider_polling', async () => providerConfig());
const result = await atStage('provider_polling', () =>
providerRequest(`${provider.base}/images/edits/${encodeURIComponent(taskId)}`, provider.key));
const status = String(result.status ?? '');
if (['failed', 'error', 'cancelled'].includes(status))
throw new PcoreTerminalError(pcoreFailureCode(result.error),
JSON.stringify(result.error ?? result));
if (!['completed', 'succeeded'].includes(status)) return null;
const image = await atStage('result_download', async () => {
try { return providerImageResult(result); }
catch (error) {
throw new PcoreResponseError(error instanceof Error ? error.message : 'Pcore 图片结果无效', null);
}
});
return typeof image === 'string'
? atStage('result_download', () => downloadProviderImage(image)) : image;
}
async function schedulePoll(database: Database, workspaceId: string, id: string) {
await database.workspaceTransaction(workspaceId, async (tx) => {
const current = await tx.query<{ status: string }>(
'SELECT status FROM atlas.media_ai_edits WHERE id = $1 FOR UPDATE', [id],
);
if (current.rows[0]?.status !== 'running') return;
await tx.query(
`INSERT INTO atlas.outbox_events
(workspace_id, aggregate_type, aggregate_id, event_type, event_payload, deduplication_key, next_attempt_at)
VALUES (current_setting('app.workspace_id')::uuid, 'media_ai_edit', $1,
'media.ai_edit.requested', jsonb_build_object('edit_id', $1::uuid), $2, now() + interval '3 seconds')`,
[id, `${id}:${randomUUID()}`],
);
});
}
export async function normalizeAiEditOutput(bytes: Buffer, settings: MediaAiEditSettings, source?: Buffer): Promise<Buffer> {
const target = sizes[settings.mode === 'outpaint' ? '3:4' : settings.ratio];
const image = sharp(bytes, { limitInputPixels: 100_000_000, failOn: 'error' });
const meta = await image.metadata();
if (!['png', 'jpeg', 'webp'].includes(meta.format ?? '') || !meta.width || !meta.height)
throw new Error('Pcore 结果图片无效');
if (settings.mode === 'outpaint' && source) {
const restored = await outpaintInputs(source);
return image.rotate().resize(target[0], target[1], { fit: 'cover' })
.composite([{ input: restored.original, left: restored.left, top: restored.top }])
.flatten({ background: '#ffffff' }).jpeg({ quality: 92 }).toBuffer();
}
return image.rotate().resize(target[0], target[1], { fit: 'contain', background: '#ffffff' })
.flatten({ background: '#ffffff' }).jpeg({ quality: 92 }).toBuffer();
}
/** 每个 Outbox 事件只推进一次供应商状态,等待由下一条延迟事件承担。 */
export async function runMediaAiEdit(database: Database, storage: S3Client, job: OutboxJob) {
const id = String(job.payload.edit_id ?? '');
const task = await atStage('task_setup', () => database.workspaceTransaction(job.workspaceId, async (tx) => {
const result = await tx.query<EditTask>(
'SELECT status, source_asset_id, reference_asset_ids, settings, provider_task_id, created_at FROM atlas.media_ai_edits WHERE id = $1 FOR UPDATE', [id],
);
const current = result.rows[0];
if (!current || !['queued', 'running'].includes(current.status)) return null;
if (Date.now() - new Date(current.created_at).getTime() > 20 * 60_000)
throw new AiEditStageError('provider_polling', new PcoreResponseError('等待 Pcore 图片任务超时', null));
await tx.query(
`UPDATE atlas.media_ai_edits SET status = 'running', started_at = coalesce(started_at, now()) WHERE id = $1`, [id],
);
await tx.query(
`UPDATE atlas.operations SET status = 'running', started_at = coalesce(started_at, now()),
attempt_count = attempt_count + 1 WHERE id = $1 AND status IN ('queued', 'running')`, [id],
);
return current;
}));
if (!task) return;
const assetIds = [task.source_asset_id, ...task.reference_asset_ids];
const assets = await atStage('input_preparation', () => database.workspaceTransaction(job.workspaceId, async (tx) => {
const result = await tx.query<StoredAsset>(
`SELECT id, storage_provider, bucket, object_key FROM atlas.media_assets WHERE id = ANY($1::uuid[])`, [assetIds],
);
const byId = new Map(result.rows.map((item) => [item.id, item]));
return assetIds.map((assetId) => {
const asset = byId.get(assetId);
if (!asset) throw new Error('参考图片素材不存在');
return asset;
});
}));
if (!task.provider_task_id) {
const providerTaskId = await submitProvider(storage, assets, task.settings);
await atStage('task_setup', () => database.workspaceTransaction(job.workspaceId, async (tx) => {
await tx.query(
`UPDATE atlas.media_ai_edits SET provider_task_id = $2 WHERE id = $1 AND status = 'running' AND provider_task_id IS NULL`,
[id, providerTaskId],
);
}));
await atStage('task_setup', () => schedulePoll(database, job.workspaceId, id));
return;
}
let result: Buffer | null;
try {
result = await pollProvider(task.provider_task_id);
} catch (error) {
// 供应商任务已是终态,重试同一个任务 ID 不会改变结果。
if (!(error instanceof PcoreTerminalError)) throw error;
await markMediaAiEditFailed(database, job, {
code: error.errorCode, source: 'pcore', stage: 'provider_polling',
detail: safeFailureDetail(error, 'provider_polling'), providerErrorRaw: error.rawResponse,
});
return;
}
if (!result) {
await atStage('task_setup', () => schedulePoll(database, job.workspaceId, id));
return;
}
const source = task.settings.mode === 'outpaint'
? await atStage('result_processing', () => readObject(storage, assets[0])) : undefined;
const output = await atStage('result_processing', () => normalizeAiEditOutput(result, task.settings, source));
if (output.byteLength > 20 * 1024 * 1024)
throw new AiEditStageError('result_processing', new Error('AI 编辑结果超过素材上限'));
const checksum = createHash('sha256').update(output).digest('hex');
const objectKey = `ai-edits/${job.workspaceId}/${checksum}.jpg`;
const existing = await atStage('result_storage', () => database.workspaceTransaction(job.workspaceId, async (tx) => {
const row = await tx.query<{ id: string }>('SELECT id FROM atlas.media_assets WHERE checksum_sha256 = $1', [checksum]);
return row.rows[0]?.id ?? null;
}));
if (!existing) await atStage('result_storage', () => storage.send(new PutObjectCommand({
Bucket: assets[0].bucket, Key: objectKey, Body: output, ContentType: 'image/jpeg',
})));
const [width, height] = sizes[task.settings.mode === 'outpaint' ? '3:4' : task.settings.ratio];
await atStage('result_storage', () => database.workspaceTransaction(job.workspaceId, async (tx) => {
const current = await tx.query<{ status: string }>(
'SELECT status FROM atlas.media_ai_edits WHERE id = $1 FOR UPDATE', [id],
);
if (current.rows[0]?.status !== 'running') return;
const asset = await tx.query<{ id: string }>(
`INSERT INTO atlas.media_assets
(workspace_id, storage_provider, bucket, object_key, content_type, byte_size, checksum_sha256, width_px, height_px)
VALUES (current_setting('app.workspace_id')::uuid, $1, $2, $3, 'image/jpeg', $4, $5, $6, $7)
ON CONFLICT (workspace_id, checksum_sha256)
DO UPDATE SET checksum_sha256 = EXCLUDED.checksum_sha256 RETURNING id`,
[assets[0].storage_provider, assets[0].bucket, objectKey, output.byteLength, checksum, width, height],
);
const assetId = asset.rows[0].id;
await tx.query(
`UPDATE atlas.media_ai_edits SET status = 'succeeded', output_asset_id = $2, error_code = NULL,
completed_at = now() WHERE id = $1`, [id, assetId],
);
await tx.query(
`UPDATE atlas.operations SET status = 'succeeded', completed_at = now(),
result = jsonb_build_object('output_asset_id', $2::uuid), error_code = NULL, error_message = NULL
WHERE id = $1 AND status IN ('queued', 'running')`, [id, assetId],
);
}));
}
export async function markMediaAiEditFailed(
database: Database, job: OutboxJob, failure: AiEditFailure,
) {
const id = String(job.payload.edit_id ?? '');
const { code, source, stage, detail, providerErrorRaw } = failure;
const message = code === 'AI_IMAGE_CONTENT_REJECTED'
? 'Pcore 未能完成这次生图,可能是输入或生成结果被审核拦截'
: code === 'AI_IMAGE_PROVIDER_FAILED'
? 'Pcore 图片任务失败,请稍后重试'
: '单张 AI 编辑失败,请重试';
await database.workspaceTransaction(job.workspaceId, async (tx) => {
await tx.query(
`UPDATE atlas.media_ai_edits SET status = 'failed', error_code = $2,
completed_at = now() WHERE id = $1 AND status IN ('queued', 'running')`, [id, code],
);
await tx.query(
`UPDATE atlas.operations SET status = 'failed', error_code = $2,
error_message = $3,
result = coalesce(result, '{}'::jsonb) || jsonb_strip_nulls(jsonb_build_object(
'failure_source', $4::text, 'failure_stage', $5::text,
'failure_detail', $6::text, 'provider_error_raw', $7::text)),
completed_at = now()
WHERE id = $1 AND status IN ('queued', 'running')`,
[id, code, message, source, stage, detail, providerErrorRaw],
);
});
}