316 lines
10 KiB
TypeScript
316 lines
10 KiB
TypeScript
import "server-only";
|
|
|
|
import { createHash } from "node:crypto";
|
|
import {
|
|
isGeneratedImageSource,
|
|
collectImageArchiveSources,
|
|
replaceArchivedImageUrls,
|
|
} from "./image-archive-urls";
|
|
import { and, asc, eq, inArray, lt, lte, or, sql } from "drizzle-orm";
|
|
import { db } from "@/lib/db";
|
|
import { imageArchiveTask, type ImageArchiveTask } from "@/lib/db/schema";
|
|
|
|
export type ImageArchiveTaskRecord = ImageArchiveTask;
|
|
|
|
const MAX_IMAGE_ARCHIVE_ATTEMPTS = 12;
|
|
const ARCHIVE_LEASE_MS = 30 * 60 * 1000;
|
|
|
|
// Replace as soon as the durable object exists, regardless of legacy hold settings.
|
|
export function getImageArchiveDisplayHoldMs() {
|
|
return 0;
|
|
}
|
|
|
|
export function imageArchiveSourceKey(source: string) {
|
|
if (!source.startsWith("data:")) return source;
|
|
const configured =
|
|
process.env.NEXT_PUBLIC_APP_URL || process.env.AUTH_URL || process.env.NEXTAUTH_URL;
|
|
const origin = configured
|
|
? new URL(configured).origin
|
|
: `http://localhost:${process.env.PORT || "3000"}`;
|
|
const key = createHash("sha256").update(source).digest("hex");
|
|
return `${origin}/api/generated-image-preview/${key}`;
|
|
}
|
|
|
|
export async function enqueueImageArchiveTask(params: {
|
|
sourceUrl: string;
|
|
archiveAfterAt?: Date;
|
|
replaceAfterAt?: Date;
|
|
metadata?: Record<string, unknown> | null;
|
|
}) {
|
|
const sourceUrl = String(params.sourceUrl || "").trim();
|
|
if (!isGeneratedImageSource(sourceUrl)) return null;
|
|
|
|
const now = new Date();
|
|
const archiveAfterAt = params.archiveAfterAt || now;
|
|
const replaceAfterAt =
|
|
params.replaceAfterAt || new Date(now.getTime() + getImageArchiveDisplayHoldMs());
|
|
|
|
const [task] = await db
|
|
.insert(imageArchiveTask)
|
|
.values({
|
|
sourceUrl: imageArchiveSourceKey(sourceUrl),
|
|
status: "pending",
|
|
archiveAfterAt,
|
|
replaceAfterAt,
|
|
metadata: sourceUrl.startsWith("data:")
|
|
? { ...params.metadata, inlineSource: sourceUrl }
|
|
: params.metadata || null,
|
|
updatedAt: now,
|
|
})
|
|
.onConflictDoUpdate({
|
|
target: imageArchiveTask.sourceUrl,
|
|
set: {
|
|
replaceAfterAt,
|
|
metadata: sql`case when ${imageArchiveTask.targetUrl} is not null then ${imageArchiveTask.metadata}
|
|
else coalesce(${imageArchiveTask.metadata}, '{}'::jsonb) || ${JSON.stringify(
|
|
sourceUrl.startsWith("data:")
|
|
? { ...params.metadata, inlineSource: sourceUrl }
|
|
: params.metadata || {},
|
|
)}::jsonb end`,
|
|
},
|
|
})
|
|
.returning();
|
|
|
|
return task || null;
|
|
}
|
|
|
|
export async function enqueueImageArchiveTasks(params: {
|
|
urls: string[];
|
|
metadata?: Record<string, unknown> | null;
|
|
}) {
|
|
const urls = Array.from(
|
|
new Set(
|
|
(params.urls || []).map((url) => String(url || "").trim()).filter(isGeneratedImageSource),
|
|
),
|
|
);
|
|
|
|
const tasks: ImageArchiveTaskRecord[] = [];
|
|
for (const sourceUrl of urls) {
|
|
const task = await enqueueImageArchiveTask({
|
|
sourceUrl,
|
|
metadata: params.metadata || null,
|
|
});
|
|
if (task) tasks.push(task);
|
|
}
|
|
return tasks;
|
|
}
|
|
|
|
export async function getImageArchiveTasksToStore(params: { limit: number }) {
|
|
const now = new Date();
|
|
return await db
|
|
.select()
|
|
.from(imageArchiveTask)
|
|
.where(
|
|
and(
|
|
inArray(imageArchiveTask.status, ["pending", "failed"]),
|
|
lte(imageArchiveTask.archiveAfterAt, now),
|
|
lt(imageArchiveTask.attemptCount, MAX_IMAGE_ARCHIVE_ATTEMPTS),
|
|
),
|
|
)
|
|
.orderBy(asc(imageArchiveTask.createdAt), asc(imageArchiveTask.id))
|
|
.limit(Math.max(1, Math.min(50, params.limit)));
|
|
}
|
|
|
|
export async function markImageArchiveTaskArchiving(id: string) {
|
|
const claimed = await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: "archiving",
|
|
attemptCount: sql`${imageArchiveTask.attemptCount} + 1`,
|
|
lastError: null,
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(
|
|
and(
|
|
eq(imageArchiveTask.id, id),
|
|
inArray(imageArchiveTask.status, ["pending", "failed"]),
|
|
lte(imageArchiveTask.archiveAfterAt, new Date()),
|
|
),
|
|
)
|
|
.returning({ id: imageArchiveTask.id });
|
|
return claimed.length > 0;
|
|
}
|
|
|
|
export async function markImageArchiveTaskStored(params: {
|
|
id: string;
|
|
targetUrl: string;
|
|
targetStorageConfigId?: string | null;
|
|
}) {
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: "stored",
|
|
targetUrl: params.targetUrl,
|
|
targetStorageConfigId: params.targetStorageConfigId || null,
|
|
archivedAt: new Date(),
|
|
replaceAfterAt: new Date(),
|
|
lastError: null,
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(eq(imageArchiveTask.id, params.id));
|
|
}
|
|
|
|
export async function markImageArchiveTaskFailed(params: { id: string; error: unknown }) {
|
|
const message = String(
|
|
params.error instanceof Error ? params.error.message : params.error || "Unknown error",
|
|
).slice(0, 2000);
|
|
const isPermanentSourceError = /\bHTTP (?:404|410)\b/i.test(message);
|
|
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: isPermanentSourceError
|
|
? "abandoned"
|
|
: sql`case
|
|
when ${imageArchiveTask.attemptCount} >= ${MAX_IMAGE_ARCHIVE_ATTEMPTS}
|
|
then 'abandoned'
|
|
else 'failed'
|
|
end`,
|
|
lastError: message,
|
|
archiveAfterAt: new Date(Date.now() + 10 * 60 * 1000),
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(eq(imageArchiveTask.id, params.id));
|
|
}
|
|
|
|
function isArchiveReplacementDue(now: Date) {
|
|
return and(
|
|
or(
|
|
eq(imageArchiveTask.status, "stored"),
|
|
and(
|
|
eq(imageArchiveTask.status, "completed"),
|
|
sql`${imageArchiveTask.metadata}->>'replacementPasses' = '1'`,
|
|
),
|
|
),
|
|
lte(imageArchiveTask.replaceAfterAt, now),
|
|
);
|
|
}
|
|
|
|
export async function getImageArchiveTasksToReplace(params: { limit: number }) {
|
|
return await db
|
|
.select()
|
|
.from(imageArchiveTask)
|
|
.where(isArchiveReplacementDue(new Date()))
|
|
.orderBy(asc(imageArchiveTask.createdAt), asc(imageArchiveTask.id))
|
|
.limit(Math.max(1, Math.min(50, params.limit)));
|
|
}
|
|
|
|
export async function markImageArchiveTaskReplacing(id: string) {
|
|
const claimed = await db
|
|
.update(imageArchiveTask)
|
|
.set({ status: "replacing", updatedAt: new Date() })
|
|
.where(and(eq(imageArchiveTask.id, id), isArchiveReplacementDue(new Date())))
|
|
.returning({ id: imageArchiveTask.id });
|
|
return claimed.length > 0;
|
|
}
|
|
|
|
export async function markImageArchiveTaskCompleted(params: {
|
|
id: string;
|
|
replacedCount?: number;
|
|
}) {
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: "completed",
|
|
replacedAt: new Date(),
|
|
lastError: null,
|
|
// Reconcile once more after delivery has settled: asynchronous generation/task
|
|
// writes may arrive just after the first replacement. The live URL is already ready.
|
|
replaceAfterAt: new Date(Date.now() + 60_000),
|
|
metadata: sql`(coalesce(${imageArchiveTask.metadata}, '{}'::jsonb) - 'inlineSource')
|
|
|| jsonb_build_object('replacementPasses', coalesce((${imageArchiveTask.metadata}->>'replacementPasses')::int, 0) + 1)
|
|
|| ${JSON.stringify(params.replacedCount === undefined ? {} : { replacedCount: params.replacedCount })}::jsonb`,
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(eq(imageArchiveTask.id, params.id));
|
|
}
|
|
|
|
export async function markImageArchiveTaskReplaceFailed(params: { id: string; error: unknown }) {
|
|
const message = String(
|
|
params.error instanceof Error ? params.error.message : params.error || "Unknown error",
|
|
).slice(0, 2000);
|
|
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: "stored",
|
|
lastError: message,
|
|
replaceAfterAt: new Date(Date.now() + 10 * 60 * 1000),
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(eq(imageArchiveTask.id, params.id));
|
|
}
|
|
|
|
export async function getImageArchiveTaskStats() {
|
|
const grouped = await db
|
|
.select({
|
|
status: imageArchiveTask.status,
|
|
count: sql<number>`count(*)::int`,
|
|
reconciliationCount: sql<number>`count(*) filter (where ${imageArchiveTask.status} = 'completed' and ${imageArchiveTask.metadata}->>'replacementPasses' = '1')::int`,
|
|
})
|
|
.from(imageArchiveTask)
|
|
.groupBy(imageArchiveTask.status);
|
|
|
|
const pendingStore = grouped
|
|
.filter((item) => item.status === "pending" || item.status === "failed")
|
|
.reduce((sum, item) => sum + (item.count || 0), 0);
|
|
const pendingReplace =
|
|
grouped
|
|
.filter((item) => item.status === "stored")
|
|
.reduce((sum, item) => sum + (item.count || 0), 0) +
|
|
grouped.reduce((sum, item) => sum + (item.reconciliationCount || 0), 0);
|
|
|
|
return {
|
|
grouped,
|
|
pendingStore,
|
|
pendingReplace,
|
|
};
|
|
}
|
|
|
|
// Interrupted workers can be claimed again after a lease expires. Failed uploads keep
|
|
// their original source and retry time, so transient bucket errors cannot lose images.
|
|
export async function recoverStaleImageArchiveTasks() {
|
|
const staleBefore = new Date(Date.now() - ARCHIVE_LEASE_MS);
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({
|
|
status: sql`case when ${imageArchiveTask.attemptCount} >= ${MAX_IMAGE_ARCHIVE_ATTEMPTS} then 'abandoned' else 'failed' end`,
|
|
archiveAfterAt: new Date(),
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(
|
|
and(eq(imageArchiveTask.status, "archiving"), lt(imageArchiveTask.updatedAt, staleBefore)),
|
|
);
|
|
await db
|
|
.update(imageArchiveTask)
|
|
.set({ status: "stored", replaceAfterAt: new Date(), updatedAt: new Date() })
|
|
.where(
|
|
and(eq(imageArchiveTask.status, "replacing"), lt(imageArchiveTask.updatedAt, staleBefore)),
|
|
);
|
|
}
|
|
|
|
export async function getArchivedImageUrlMappings(urls: string[]) {
|
|
if (!urls.length) return {};
|
|
const sources = [...new Set(urls.filter(isGeneratedImageSource))];
|
|
const tasks = await db
|
|
.select({ sourceUrl: imageArchiveTask.sourceUrl, targetUrl: imageArchiveTask.targetUrl })
|
|
.from(imageArchiveTask)
|
|
.where(inArray(imageArchiveTask.sourceUrl, sources.map(imageArchiveSourceKey)));
|
|
const targets = new Map(
|
|
tasks.filter((task) => task.targetUrl).map((task) => [task.sourceUrl, task.targetUrl!]),
|
|
);
|
|
const replacements: Record<string, string> = {};
|
|
for (const source of sources) {
|
|
const target = targets.get(imageArchiveSourceKey(source));
|
|
if (target) replacements[source] = target;
|
|
}
|
|
return replacements;
|
|
}
|
|
|
|
// Normalize late writes too: an archive can finish before a task/canvas save arrives.
|
|
export async function resolveArchivedImageUrls<T>(value: T): Promise<T> {
|
|
const sources = collectImageArchiveSources(value);
|
|
if (!sources.length) return value;
|
|
return replaceArchivedImageUrls(value, await getArchivedImageUrlMappings(sources));
|
|
}
|