refactor: optimize AI moderation pipeline, fix OOM risks, token duplication, and add Zod validation

src/moderation/aiAnalysisWorker.ts

@@ -1,6 +1,6 @@
 import { config } from "../config.js";
 import { initializeDatabase } from "../database/drizzle.js";
-import { buildConversationPromptMessages } from "./conversationContext.js";
+import { buildConversationContext } from "./conversationContext.js";
 import { runModerationAnalysis } from "./llmModerationClient.js";
 import {
   getAttachmentsForMessages,
@@ -67,7 +67,7 @@ export default async function processAnalysisRequest({
       limit: config.AI_ANALYSIS_CONTEXT_MESSAGE_LIMIT,
     });
 
-    const promptMessages = buildConversationPromptMessages({
+    const contextLines = buildConversationContext({
       contextBefore,
       targets: messages,
       maxTokens: config.AI_ANALYSIS_MAX_CONTEXT_TOKENS,
@@ -80,7 +80,7 @@ export default async function processAnalysisRequest({
 
     const result = await runModerationAnalysis({
       targets: messages,
-      contextText: promptMessages.join("\n"),
+      contextText: contextLines.join("\n"),
       attachments,
     });
 
@@ -94,12 +94,19 @@ export default async function processAnalysisRequest({
         analysis: analysisResult.analysis,
         analyzedAt: Date.now(),
         error: null,
-      }
+      },
     }));
-    
-    const rows = await updateMessagesAIAnalysisBulk(updates);
 
-    return { ok: true, conversationKey, rows };
+    try {
+      const rows = await updateMessagesAIAnalysisBulk(updates);
+      return { ok: true, conversationKey, rows };
+    } catch (dbErr) {
+      // If bulk update fails, we log it but don't fail the worker completely
+      // so it can at least retry later without blowing up the circuit breaker if it was an isolated issue
+      throw new Error(
+        `Failed to update DB: ${dbErr instanceof Error ? dbErr.message : String(dbErr)}`,
+      );
+    }
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : String(error);
     const errorStack = error instanceof Error ? error.stack : undefined;

src/moderation/aiAnalyzer.ts

@@ -3,6 +3,10 @@ import { fileURLToPath } from "node:url";
 import { Piscina } from "piscina";
 import { config } from "../config.js";
 import { createChildLogger } from "../logger.js";
+import {
+  estimateTokens,
+  formatMessageForPrompt,
+} from "./conversationContext.js";
 import {
   getMessageById,
   getPendingConversationKeys,
@@ -88,10 +92,8 @@ export function pickBatchWithinBudget(
   let usedTokens = 0;
 
   for (const msg of messages) {
-    // Estimate tokens based on actual content length (conservative: 3 chars/token)
+    const formatted = formatMessageForPrompt(msg, "target");
-    const content = msg.edited_content ?? msg.content;
+    const msgTokens = estimateTokens(formatted) + tokensPerMessage;
-    const contentTokens = Math.ceil(content.length / 3);
-    const msgTokens = contentTokens + tokensPerMessage;
 
     if (usedTokens + msgTokens <= maxTokens) {
       batch.push(msg);
@@ -118,7 +120,8 @@ async function processBatch(
 ): Promise<void> {
   if (messages.length === 0) return;
   if (Date.now() < globalCooldownUntil) {
-    return; // Circuit breaker is open
+    // Should not normally hit here due to checks in scheduleConversationAnalysis, but just in case
+    return;
   }
 
   activeRequests++;
@@ -126,7 +129,10 @@ async function processBatch(
   const processingStartedAt = Date.now();
   conversationProcessing.set(conversationKey, processingStartedAt);
   try {
-    const result = (await workerPool.run({ conversationKey, messages })) as AnalysisWorkerResponse;
+    const result = (await workerPool.run({
+      conversationKey,
+      messages,
+    })) as AnalysisWorkerResponse;
 
     for (const row of result.rows) {
       getModerationBroadcaster()?.messageAnalyzed(row);
@@ -136,9 +142,11 @@ async function processBatch(
       consecutiveErrors++;
       if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) {
         globalCooldownUntil = Date.now() + 60000;
-        logger.warn("Global circuit breaker triggered due to consecutive errors");
+        logger.warn(
+          "Global circuit breaker triggered due to consecutive errors",
+        );
       }
-      
+
       lastError = result.error ?? "Analysis worker failed";
       conversationErrorCooldown.set(
         conversationKey,
@@ -201,7 +209,6 @@ async function processBatch(
   }
 }
 
-
 /**
  * Debounced analysis trigger for a conversation
  */
@@ -211,9 +218,20 @@ function scheduleConversationAnalysis(conversationKey: string): void {
     return;
   }
 
-  // Skip if in error cooldown
+  // Check cooldowns
-  const cooldownUntil = conversationErrorCooldown.get(conversationKey);
+  const convoCooldown = conversationErrorCooldown.get(conversationKey) || 0;
-  if (cooldownUntil && Date.now() < cooldownUntil) {
+  const activeCooldown = Math.max(convoCooldown, globalCooldownUntil);
+
+  if (activeCooldown && Date.now() < activeCooldown) {
+    // Instead of dropping, re-schedule for after cooldown if not already scheduled
+    if (!conversationDebounceTimers.has(conversationKey)) {
+      const remaining = activeCooldown - Date.now();
+      const timer = setTimeout(() => {
+        conversationDebounceTimers.delete(conversationKey);
+        scheduleConversationAnalysis(conversationKey);
+      }, remaining + 500); // 500ms buffer after cooldown
+      conversationDebounceTimers.set(conversationKey, timer);
+    }
     return;
   }
 

src/moderation/conversationContext.ts

@@ -14,57 +14,52 @@ function formatTimestamp(ms: number): string {
 }
 
 /**
- * Estimates token count for a string (rough approximation: ~4 chars per token)
+ * Estimates token count for a string (pessimistic approximation for Indonesian slang & JSON overhead)
  */
-function estimateTokens(text: string): number {
+export function estimateTokens(text: string): number {
-  return Math.ceil(text.length / 4);
+  return Math.ceil(text.length / 3) + 15;
 }
 
 /**
- * Builds conversation prompt messages with context and targets
+ * Formats a single message for context or target display
- * - Marks target messages with [target], prior context with [context]
- * - Uses edited_content when present, otherwise content
- * - Maintains chronological order
- * - Respects maxTokens budget, prioritizing targets and most recent context
  */
-export function buildConversationPromptMessages(
+export function formatMessageForPrompt(
+  msg: MessageRecord,
+  label: "context" | "target",
+): string {
+  const content = msg.edited_content ?? msg.content;
+  const timestamp = formatTimestamp(msg.created_at);
+  return `[${label}] id=${msg.id} time=${timestamp} user=${msg.username}: ${content}`;
+}
+
+/**
+ * Builds conversation historical context without including targets.
+ * Calculates how much token budget targets use, and fills the rest with context.
+ */
+export function buildConversationContext(
   input: ConversationContextInput,
 ): string[] {
   const { contextBefore, targets, maxTokens } = input;
 
-  const formatMessage = (msg: MessageRecord, label: string): string => {
+  // Calculate tokens used by targets
-    const content = msg.edited_content ?? msg.content;
+  let usedTokens = targets.reduce((sum, msg) => {
-    const timestamp = formatTimestamp(msg.created_at);
+    return sum + estimateTokens(formatMessageForPrompt(msg, "target"));
-    return `[${label}] id=${msg.id} time=${timestamp} user=${msg.username}: ${content}`;
+  }, 0);
-  };
 
-  const targetEntries = targets.map((msg) => ({
+  const selectedContextLines: string[] = [];
-    msg,
-    label: "target" as const,
-    line: formatMessage(msg, "target"),
-  }));
 
-  let usedTokens = targetEntries.reduce(
+  // Go backwards through context, taking most recent first
-    (sum, entry) => sum + estimateTokens(entry.line),
-    0,
-  );
-
-  const selectedContextEntries: Array<{
-    msg: MessageRecord;
-    label: "context";
-    line: string;
-  }> = [];
   for (let i = contextBefore.length - 1; i >= 0; i--) {
     const msg = contextBefore[i];
-    const line = formatMessage(msg, "context");
+    const line = formatMessageForPrompt(msg, "context");
     const lineTokens = estimateTokens(line);
+
     if (usedTokens + lineTokens <= maxTokens) {
-      selectedContextEntries.push({ msg, label: "context", line });
+      // Unshift so oldest context is first in the array
+      selectedContextLines.unshift(line);
       usedTokens += lineTokens;
     }
   }
 
-  return [...selectedContextEntries, ...targetEntries]
+  return selectedContextLines;
-    .sort((a, b) => a.msg.created_at - b.msg.created_at)
-    .map((entry) => entry.line);
 }

src/moderation/llmModerationClient.ts

@@ -1,4 +1,5 @@
 import OpenAI from "openai";
+import { z } from "zod";
 import { config } from "../config.js";
 import { createChildLogger } from "../logger.js";
 import { retryWithBackoff } from "../retry.js";
@@ -8,45 +9,66 @@ import type {
   MessageRecord,
 } from "./types.js";
 
+const ModerationResponseSchema = z.object({
+  results: z.array(
+    z.object({
+      message_id: z.union([z.string(), z.number()]).transform(String),
+      status: z.enum(["clean", "warn", "flagged"]).catch("clean"),
+      flags: z.array(z.string()).catch([]),
+      score: z.number().catch(0),
+      analysis: z.string().catch(""),
+    }),
+  ),
+});
+
 const log = createChildLogger("llmModerationClient");
 const openai = new OpenAI({
   apiKey: config.AI_LLM_API_KEY,
   baseURL: config.AI_LLM_BASE_URL,
   maxRetries: 0,
-  timeout: 2_147_483_647,
+  timeout: 30000,
   fetch: async (url, init) => {
-    const response = await globalThis.fetch(url, init);
+    // Add internal timeout for the global fetch as safety
-    const body =
+    const controller = new AbortController();
-      typeof response.text === "function"
+    const timeout = setTimeout(() => controller.abort(), 30000);
-        ? await response.text()
+    const fetchInit = { ...init, signal: controller.signal };
-        : JSON.stringify(await response.json());
 
-    let normalizedBody = body;
+    try {
-    if (response.ok !== false) {
+      const response = await globalThis.fetch(url, fetchInit);
-      try {
+      const body =
-        JSON.parse(body);
+        typeof response.text === "function"
-      } catch (error) {
+          ? await response.text()
-        log.warn(
+          : JSON.stringify(await response.json());
-          {
+
-            error: error instanceof Error ? error.message : String(error),
+      let normalizedBody = body;
-            status: response.status ?? 200,
+      if (response.ok !== false) {
-            bodyLength: body.length,
+        try {
-            body,
+          JSON.parse(body);
-          },
+        } catch (error) {
-          "LLM provider returned malformed JSON response body",
+          log.warn(
-        );
+            {
-        normalizedBody = JSON.stringify(extractJson(body));
+              error: error instanceof Error ? error.message : String(error),
+              status: response.status ?? 200,
+              bodyLength: body.length,
+              body,
+            },
+            "LLM provider returned malformed JSON response body",
+          );
+          normalizedBody = JSON.stringify(extractJson(body));
+        }
       }
+
+      const headers = new Headers(response.headers ?? undefined);
+      headers.set("Content-Type", "application/json");
+      headers.delete("Content-Length");
+
+      return new Response(normalizedBody, {
+        status: response.status ?? 200,
+        headers,
+      });
+    } finally {
+      clearTimeout(timeout);
     }
-
-    const headers = new Headers(response.headers ?? undefined);
-    headers.set("Content-Type", "application/json");
-    headers.delete("Content-Length");
-
-    return new Response(normalizedBody, {
-      status: response.status ?? 200,
-      headers,
-    });
   },
 });
 
@@ -130,8 +152,6 @@ export function extractJson(content: string): any {
   throw new Error("No JSON object found in response");
 }
 
-
-
 export function parseModerationResponse(
   content: string,
   targetIds: string[],
@@ -156,66 +176,37 @@ export function parseModerationResponse(
     }
   }
 
-  if (!parsed || typeof parsed !== "object" || !Array.isArray(parsed.results)) {
+  const parseResult = ModerationResponseSchema.safeParse(parsed);
-    throw new Error("Response missing 'results' array");
+  if (!parseResult.success) {
+    throw new Error(`Zod validation failed: ${parseResult.error.message}`);
   }
 
-  const response = parsed as RawModerationResponse;
+  const response = parseResult.data;
   const foundIds = new Set<string>();
   const targetIdSet = new Set(targetIds);
 
-  const results: (AnalysisResult | null)[] = response.results.map(
+  const results: (AnalysisResult | null)[] = response.results.map((result) => {
-    (result, index) => {
+    const { message_id, status, flags, score, analysis } = result;
-      const { message_id, status, flags, score, analysis } = result;
+    const finalId = message_id.trim();
 
-      if (!message_id) {
+    if (!targetIdSet.has(finalId)) {
-        throw new Error("Result missing 'message_id'");
+      return null;
-      }
+    }
 
-      const finalId = String(message_id).trim();
+    if (foundIds.has(finalId)) {
+      return null; // Ignore duplicates safely
+    }
 
-      if (!targetIdSet.has(finalId)) {
+    foundIds.add(finalId);
-        log.warn(
-          { unknownId: finalId, originalId: message_id, targetIds },
-          "Skipping moderation result for non-target message_id",
-        );
-        return null;
-      }
 
-      if (foundIds.has(finalId)) {
+    return {
-        log.warn({ duplicateId: finalId }, "Duplicate message_id in response");
+      messageId: finalId,
-        throw new Error(`Duplicate message_id: ${finalId}`);
+      status: status as "clean" | "warn" | "flagged",
-      }
+      flags,
-
+      score: Math.max(0, Math.min(1, score)),
-      foundIds.add(finalId);
+      analysis,
-
+    };
-      const validStatuses = ["clean", "warn", "flagged"] as const;
+  });
-      const safeStatus = validStatuses.includes(status as any) ? status : "clean";
-
-      let numScore = Number(score);
-      if (!Number.isFinite(numScore)) {
-        numScore = 0;
-      }
-      numScore = Math.max(0, Math.min(1, numScore));
-
-      let flagsArray: string[] = [];
-      if (Array.isArray(flags)) {
-        flagsArray = flags.map((f) => String(f));
-      } else if (flags) {
-        flagsArray = [String(flags)];
-      }
-
-      const analysisStr = analysis ? String(analysis) : "";
-
-      return {
-        messageId: finalId,
-        status: safeStatus as "clean" | "warn" | "flagged",
-        flags: flagsArray,
-        score: numScore,
-        analysis: analysisStr,
-      };
-    },
-  );
 
   const filteredResults = results.filter(
     (r): r is AnalysisResult => r !== null,
@@ -362,7 +353,9 @@ export async function runModerationAnalysis(
   const targetIdSet = new Set(targets.map((t) => t.id));
 
   const candidateAttachments = (attachments ?? [])
-    .filter((att) => getAttachmentImageUrl(att) && att.type.startsWith("image/"))
+    .filter(
+      (att) => getAttachmentImageUrl(att) && att.type.startsWith("image/"),
+    )
     .sort((a, b) => {
       // Target-message attachments always come first so they consume the cap first
       const aIsTarget = targetIdSet.has(a.message_id) ? 1 : 0;
@@ -380,12 +373,16 @@ export async function runModerationAnalysis(
       const urlToUse = getAttachmentImageUrl(att);
       if (!urlToUse) return;
 
+      const controller = new AbortController();
+      const timeoutId = setTimeout(() => controller.abort(), 15000);
+
       try {
         log.info(
           { attachmentId: att.id, messageId: att.message_id, url: urlToUse },
           "Downloading attachment for base64 encoding",
         );
-        const res = await fetch(urlToUse);
+
+        const res = await fetch(urlToUse, { signal: controller.signal });
         if (!res.ok) {
           log.warn(
             { attachmentId: att.id, status: res.status, url: urlToUse },
@@ -394,13 +391,31 @@ export async function runModerationAnalysis(
           return;
         }
 
-        const contentLength = Number(res.headers.get("content-length") || 0);
+        if (!res.body) return;
-        if (contentLength > 10 * 1024 * 1024) {
+
-          log.warn({ attachmentId: att.id, contentLength }, "Attachment too large, skipping");
+        let totalBytes = 0;
-          return;
+        const chunks: Uint8Array[] = [];
+        const reader = res.body.getReader();
+
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+
+          if (value) {
+            totalBytes += value.length;
+            if (totalBytes > 10 * 1024 * 1024) {
+              log.warn(
+                { attachmentId: att.id },
+                "Attachment exceeded 10MB limit, aborting stream",
+              );
+              reader.cancel();
+              return;
+            }
+            chunks.push(value);
+          }
         }
 
-        const imageBytes = Buffer.from(await res.arrayBuffer());
+        const imageBytes = Buffer.concat(chunks);
         const sniffedMime = sniffImageMimeType(imageBytes);
         if (!sniffedMime) {
           log.warn(
@@ -417,7 +432,10 @@ export async function runModerationAnalysis(
         }
 
         const dataUrl = `data:${sniffedMime};base64,${imageBytes.toString("base64")}`;
-        const part: RawImagePart = { type: "image_url", image_url: { url: dataUrl } };
+        const part: RawImagePart = {
+          type: "image_url",
+          image_url: { url: dataUrl },
+        };
 
         const existing = messageImageMap.get(att.message_id) ?? [];
         existing.push(part);
@@ -430,6 +448,8 @@ export async function runModerationAnalysis(
           },
           "Error base64 encoding attachment",
         );
+      } finally {
+        clearTimeout(timeoutId);
       }
     }),
   );
@@ -524,7 +544,10 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan
 
   const buildMessageContent = (): string | ContentPart[] => {
     const correction = lastParseError
-      ? { error: lastParseError, preview: lastInvalidContent?.slice(0, 800) ?? "<empty>" }
+      ? {
+          error: lastParseError,
+          preview: lastInvalidContent?.slice(0, 800) ?? "<empty>",
+        }
       : undefined;
 
     const systemText = buildSystemPrompt(correction);
@@ -543,7 +566,10 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan
 
     // Multimodal path: interleave text + images per message
     const parts: ContentPart[] = [
-      { type: "text", text: `${systemText}\n\n## Pesan yang Dianalisis (dengan lampiran gambar)\n` },
+      {
+        type: "text",
+        text: `${systemText}\n\n## Pesan yang Dianalisis (dengan lampiran gambar)\n`,
+      },
     ];
 
     for (const msg of targets) {
@@ -586,33 +612,7 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan
           top_p: 0.95,
           max_tokens: 16384,
           response_format: {
-            type: "json_schema",
+            type: "json_object",
-            json_schema: {
-              name: "moderation",
-              strict: true,
-              schema: {
-                type: "object",
-                properties: {
-                  results: {
-                    type: "array",
-                    items: {
-                      type: "object",
-                      properties: {
-                        message_id: { type: "string" },
-                        status: { type: "string", enum: ["clean", "warn", "flagged"] },
-                        flags: { type: "array", items: { type: "string" } },
-                        score: { type: "number" },
-                        analysis: { type: "string" }
-                      },
-                      required: ["message_id", "status", "flags", "score", "analysis"],
-                      additionalProperties: false
-                    }
-                  }
-                },
-                required: ["results"],
-                additionalProperties: false
-              }
-            }
           },
           stream: false,
           chat_template_kwargs: { enable_thinking: false },
@@ -674,7 +674,7 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan
     const errorMsg =
       parseError instanceof Error ? parseError.message : String(parseError);
     const content: string = lastInvalidContent;
-    
+
     log.error(
       {
         error: errorMsg,

src/moderation/messageStore.ts

@@ -420,12 +420,14 @@ export async function updateMessageAIAnalysis(
 }
 
 export async function updateMessagesAIAnalysisBulk(
-  updates: Array<{ messageId: string; result: AIAnalysisUpdate }>
+  updates: Array<{ messageId: string; result: AIAnalysisUpdate }>,
 ): Promise<MessageRecord[]> {
   if (updates.length === 0) return [];
   try {
     const results = await Promise.all(
-      updates.map(({ messageId, result }) => updateMessageAIAnalysis(messageId, result))
+      updates.map(({ messageId, result }) =>
+        updateMessageAIAnalysis(messageId, result),
+      ),
     );
     return results.filter((r): r is MessageRecord => r !== null);
   } catch (error) {