fix(ai-moderation): resolve race conditions and implement processing state

Introduces a `processing` state to the AI analysis lifecycle to prevent
duplicate processing of the same messages.

- Implements row-level locking using `FOR UPDATE SKIP LOCKED` in
  `messageStore.ts` to ensure atomic message acquisition.
- Adds a `processing` status to the `AIStatus` type and database schema.
- Fixes a TOCTOU race condition in `aiAnalyzer.ts` by synchronizing
  the conversation processing lock before async database operations.
- Implements `revertStuckProcessingMessages` to recover messages stuck
  in the `processing` state due to worker crashes or timeouts.
- Updates `processBatch` and scheduling logic to correctly manage and
  release conversation-level locks.
This commit is contained in:
MythEclipse
2026-06-05 16:36:14 +07:00
parent 399919ded0
commit 09f6e80ddd
4 changed files with 98 additions and 14 deletions
@@ -651,8 +651,8 @@ export async function getPendingMessagesByConversation(
// conversationKey is either thread_id or channel_id
// Query both to safely handle the key
const rows = await database
.select()
const sq = database
.select({ id: messagesTable.id })
.from(messagesTable)
.where(
and(
@@ -665,7 +665,14 @@ export async function getPendingMessagesByConversation(
),
)
.orderBy(asc(messagesTable.created_at))
.limit(limit);
.limit(limit)
.for("update", { skipLocked: true });
const rows = await database
.update(messagesTable)
.set({ ai_status: "processing", ai_analyzed_at: Date.now() })
.where(inArray(messagesTable.id, sq))
.returning();
return rows as MessageRecord[];
} catch (error) {
@@ -856,8 +863,8 @@ export async function getIncompleteMessagesByConversation(
): Promise<MessageRecord[]> {
try {
const database = db();
const rows = await database
.select()
const sq = database
.select({ id: messagesTable.id })
.from(messagesTable)
.where(
and(
@@ -874,7 +881,14 @@ export async function getIncompleteMessagesByConversation(
),
)
.orderBy(asc(messagesTable.created_at))
.limit(limit);
.limit(limit)
.for("update", { skipLocked: true });
const rows = await database
.update(messagesTable)
.set({ ai_status: "processing", ai_analyzed_at: Date.now() })
.where(inArray(messagesTable.id, sq))
.returning();
return rows as MessageRecord[];
} catch (error) {
@@ -1247,3 +1261,38 @@ export async function getExpiredMessages(
throw error;
}
}
export async function revertStuckProcessingMessages(
timeoutMs: number = 300000,
): Promise<number> {
try {
const database = db();
const cutoffTime = Date.now() - timeoutMs;
const rows = await database
.update(messagesTable)
.set({ ai_status: "pending", ai_analyzed_at: null })
.where(
and(
eq(messagesTable.ai_status, "processing"),
sql`${messagesTable.ai_analyzed_at} < ${cutoffTime}`,
),
)
.returning({ id: messagesTable.id });
if (rows.length > 0) {
logger.warn(
{ count: rows.length, messageIds: rows.map((r) => r.id) },
"Reverted stuck processing messages to pending",
);
}
return rows.length;
} catch (error) {
logger.error(
{ error: error instanceof Error ? error.message : String(error) },
"Failed to revert stuck processing messages",
);
return 0;
}
}
@@ -1,7 +1,7 @@
import type fs from "node:fs";
import type prism from "prism-media";
export type AIStatus = "pending" | "clean" | "warn" | "flagged" | "error";
export type AIStatus = "pending" | "processing" | "clean" | "warn" | "flagged" | "error";
export type AISeverity = "none" | "low" | "medium" | "high" | "critical";
export type AIRecommendedAction =
| "none"