Files
GMW/services/discord-gateway/src/modules/ai-moderation/batchBudget.ts
T
asepharyana 80daa9f045 fix(gateway): un-claim budget-overflow messages stuck in processing
getPendingMessagesByConversation() flips every fetched pending row to
'processing', then pickBatchWithinBudget() may stop early on the token
budget. The tail rows that did NOT make the batch were never un-claimed,
so they stayed 'processing' forever — the recovery worker reverted them
(120s) only for the next wave to re-claim them, an infinite loop of
stuck messages that never get analyzed (saw 22 rows, some recycled for
40+ minutes).

- add computeBudgetOverflowMessages() pure helper (batchBudget.ts)
- batchScheduler un-claims overflow rows back to 'pending' before
  dispatching the trimmed batch
- recovery-worker now reverts stuck processing unconditionally (the old
  'conversationProcessing.size > 0' guard skipped the revert when the
  in-memory lock map was empty, e.g. fresh boot — exactly when stranded
  rows from a previous process need rescuing)
- 3 regression tests for the overflow helper
2026-09-24 16:49:27 +07:00

65 lines
2.3 KiB
TypeScript

/**
* batchBudget.ts
*
* Pure batch-sizing helper extracted from batchProcessor.ts so it can be
* unit-tested without pulling in the Piscina worker pool, message store,
* or any other side-effectful import chain.
*/
import type { MessageRecord } from "../message-capture/types.js";
/** Token estimator contract (satisfied by conversationContext.estimateTokens). */
export type TokenEstimator = (text: string) => number;
/**
* Picks a batch of messages within a token budget.
* `tokensPerMessage` accounts for JSON structure overhead around each entry.
* The estimator is injected so this stays a pure function — callers in the
* batch pipeline pass the tiktoken-based estimateTokens.
*/
export function pickBatchWithinBudget(
messages: MessageRecord[],
maxTokens: number,
tokensPerMessage: number,
estimateTokens: TokenEstimator,
): MessageRecord[] {
const batch: MessageRecord[] = [];
let usedTokens = 0;
for (const msg of messages) {
const content = msg.edited_content ?? msg.content;
const msgTokens = estimateTokens(content) + tokensPerMessage;
// Stop at the first overflow instead of skipping: input is ordered
// created_at ASC, so a contiguous chronological prefix keeps the batch
// gap-free. Skipped-over messages would leave unanalyzed holes mid-
// timeline; anything past the budget is picked up by the next wave
// (processBatch always re-schedules after success).
if (usedTokens + msgTokens > maxTokens) {
break;
}
batch.push(msg);
usedTokens += msgTokens;
}
return batch;
}
/**
* Returns the messages that were fetched/claimed but did NOT make it into the
* trimmed batch (i.e. the tail past the token budget).
*
* The DB claim step flips every fetched pending row to `processing`; the batch
* trim may then stop early on the token budget. Those tail rows would stay
* stuck in `processing` forever unless the caller explicitly un-claims them —
* this helper identifies exactly which rows that is, so the caller can write
* them back to `pending` for the next wave.
*/
export function computeBudgetOverflowMessages(
claimed: MessageRecord[],
trimmed: MessageRecord[],
): MessageRecord[] {
if (trimmed.length === 0) return claimed;
const trimmedIds = new Set(trimmed.map((m) => m.id));
return claimed.filter((m) => !trimmedIds.has(m.id));
}