import "server-only"; import { and, asc, eq, lt, notInArray, sql } from "drizzle-orm"; import { executeAgentPsdTemplateRenderRequest } from "@/lib/agents/psd-render-service"; import type { AgentPsdTemplateRenderRequest } from "@/lib/agents/psd-render-task-contract"; import { db } from "@/lib/db"; import { agentPsdRenderTask } from "@/lib/db/schema"; type AgentPsdRenderTaskStatus = "pending" | "running" | "succeeded" | "failed"; interface QueueState { pending: string[]; pendingSet: Set; running: Set; draining: boolean; } const DEFAULT_CONCURRENCY = 2; const MAX_CONCURRENCY = 8; const STALE_RUNNING_MS = 15 * 60 * 1000; const queueStateKey = "__agent_psd_render_task_queue_state__"; function getQueueConcurrency(): number { const raw = Number.parseInt(process.env.AGENT_PSD_RENDER_TASK_CONCURRENCY || "", 10); if (Number.isFinite(raw) && raw > 0) { return Math.min(raw, MAX_CONCURRENCY); } return DEFAULT_CONCURRENCY; } function getQueueState(): QueueState { const globalObj = globalThis as typeof globalThis & { [queueStateKey]?: QueueState; }; if (!globalObj[queueStateKey]) { globalObj[queueStateKey] = { pending: [], pendingSet: new Set(), running: new Set(), draining: false, }; } return globalObj[queueStateKey] as QueueState; } function enqueueLocalTask(taskId: string) { const state = getQueueState(); if (!taskId || state.running.has(taskId) || state.pendingSet.has(taskId)) { return; } state.pending.push(taskId); state.pendingSet.add(taskId); } async function claimNextPendingTask(): Promise { const state = getQueueState(); while (state.pending.length > 0) { const taskId = state.pending.shift(); if (!taskId) continue; state.pendingSet.delete(taskId); if (!state.running.has(taskId)) { return taskId; } } const runningIds = Array.from(state.running); const whereCondition = runningIds.length > 0 ? and(eq(agentPsdRenderTask.status, "pending"), notInArray(agentPsdRenderTask.id, runningIds)) : eq(agentPsdRenderTask.status, "pending"); const rows = await db .select({ id: agentPsdRenderTask.id }) .from(agentPsdRenderTask) .where(whereCondition) .orderBy(asc(agentPsdRenderTask.createdAt)) .limit(1); return rows[0]?.id ?? null; } async function runSingleTask(taskId: string) { const state = getQueueState(); state.running.add(taskId); const startedAtMs = Date.now(); console.log("[PSD Render Task] starting:", { taskId }); try { const [task] = await db .update(agentPsdRenderTask) .set({ status: "running", startedAt: new Date(), completedAt: null, error: null, updatedAt: new Date(), attempts: sql`${agentPsdRenderTask.attempts} + 1`, }) .where(and(eq(agentPsdRenderTask.id, taskId), eq(agentPsdRenderTask.status, "pending"))) .returning({ id: agentPsdRenderTask.id, request: agentPsdRenderTask.request, }); if (!task) { return; } const result = await executeAgentPsdTemplateRenderRequest( task.request as AgentPsdTemplateRenderRequest ); await db .update(agentPsdRenderTask) .set({ status: "succeeded", result: result as unknown as Record, error: null, completedAt: new Date(), updatedAt: new Date(), }) .where(eq(agentPsdRenderTask.id, taskId)); console.log("[PSD Render Task] succeeded:", { taskId, durationMs: Date.now() - startedAtMs, resultUrl: result.url, }); } catch (error) { console.error("[PSD Render Task] failed:", { taskId, durationMs: Date.now() - startedAtMs, error: error instanceof Error ? error.message : String(error), stack: error instanceof Error ? error.stack : undefined, }); await db .update(agentPsdRenderTask) .set({ status: "failed", error: error instanceof Error ? error.message : String(error), completedAt: new Date(), updatedAt: new Date(), }) .where(eq(agentPsdRenderTask.id, taskId)); } finally { state.running.delete(taskId); void drainAgentPsdRenderTaskQueue(); } } export async function drainAgentPsdRenderTaskQueue() { const state = getQueueState(); if (state.draining) return; state.draining = true; try { const maxConcurrent = getQueueConcurrency(); while (state.running.size < maxConcurrent) { const taskId = await claimNextPendingTask(); if (!taskId) break; if (state.running.has(taskId)) continue; void runSingleTask(taskId); } } finally { state.draining = false; } } export async function enqueueAgentPsdRenderTask(taskId: string) { enqueueLocalTask(taskId); console.log("[PSD Render Task] enqueued:", { taskId }); await drainAgentPsdRenderTaskQueue(); } export async function createAgentPsdRenderTask(params: { userId: string; request: AgentPsdTemplateRenderRequest; }) { const [task] = await db .insert(agentPsdRenderTask) .values({ userId: params.userId, status: "pending", request: params.request as unknown as Record, attempts: 0, createdAt: new Date(), updatedAt: new Date(), }) .returning({ id: agentPsdRenderTask.id, status: agentPsdRenderTask.status, createdAt: agentPsdRenderTask.createdAt, }); console.log("[PSD Render Task] created:", { taskId: task.id, userId: params.userId, templateId: params.request.templateId, }); return task; } export async function getAgentPsdRenderTaskForUser(params: { taskId: string; userId: string; }) { const rows = await db .select({ id: agentPsdRenderTask.id, status: agentPsdRenderTask.status, request: agentPsdRenderTask.request, result: agentPsdRenderTask.result, error: agentPsdRenderTask.error, attempts: agentPsdRenderTask.attempts, startedAt: agentPsdRenderTask.startedAt, completedAt: agentPsdRenderTask.completedAt, createdAt: agentPsdRenderTask.createdAt, updatedAt: agentPsdRenderTask.updatedAt, }) .from(agentPsdRenderTask) .where( and(eq(agentPsdRenderTask.id, params.taskId), eq(agentPsdRenderTask.userId, params.userId)) ) .limit(1); const task = rows[0]; if (!task) { return null; } if (task.status === "running" && task.startedAt) { const staleBefore = new Date(Date.now() - STALE_RUNNING_MS); if (task.startedAt < staleBefore) { await db .update(agentPsdRenderTask) .set({ status: "pending", startedAt: null, completedAt: null, updatedAt: new Date(), }) .where( and( eq(agentPsdRenderTask.id, task.id), eq(agentPsdRenderTask.status, "running"), lt(agentPsdRenderTask.startedAt, staleBefore) ) ); enqueueLocalTask(task.id); await drainAgentPsdRenderTaskQueue(); task.status = "pending" as AgentPsdRenderTaskStatus; task.startedAt = null; task.completedAt = null; } } if (task.status === "pending") { enqueueLocalTask(task.id); await drainAgentPsdRenderTaskQueue(); } return task; }