ai-chatbot/lib/agents/psd-render-task-queue.ts

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<string>;
  running: Set<string>;
  draining: boolean;
}

const DEFAULT_CONCURRENCY = 1;
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, 3);
  }
  return DEFAULT_CONCURRENCY;
}

function getQueueState(): QueueState {
  const globalObj = globalThis as typeof globalThis & {
    [queueStateKey]?: QueueState;
  };

  if (!globalObj[queueStateKey]) {
    globalObj[queueStateKey] = {
      pending: [],
      pendingSet: new Set<string>(),
      running: new Set<string>(),
      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<string | null> {
  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);
  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<string, unknown>,
        error: null,
        completedAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(agentPsdRenderTask.id, taskId));
    console.log("[PSD Render Task] succeeded:", {
      taskId,
      resultUrl: result.url,
    });
  } catch (error) {
    console.error("[PSD Render Task] failed:", {
      taskId,
      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<string, unknown>,
      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;
}