915 lines
27 KiB
TypeScript
915 lines
27 KiB
TypeScript
import { v4 as uuidv4 } from "uuid";
|
|
import type { BatchItemCheckpoint, BatchRecord } from "@/lib/localDb";
|
|
import {
|
|
countBatchItemCheckpoints,
|
|
createFile,
|
|
deleteFile,
|
|
ensureBatchItemCheckpoints,
|
|
getApiKeyById,
|
|
getBatch,
|
|
getFileContent,
|
|
getPendingBatches,
|
|
getTerminalBatches,
|
|
listBatchItemCheckpoints,
|
|
listFiles,
|
|
markBatchItemError,
|
|
markBatchItemProcessing,
|
|
markBatchItemResult,
|
|
updateBatch,
|
|
} from "@/lib/localDb";
|
|
import { dispatch } from "@/lib/batches/dispatch";
|
|
import type { SupportedBatchEndpoint } from "@/shared/constants/batchEndpoints";
|
|
import { DEFAULT_BATCH_EXPIRATION_SECONDS } from "@/shared/constants/batch";
|
|
|
|
let isProcessing: boolean = false;
|
|
let pollInterval: NodeJS.Timeout | null = null;
|
|
const activeProcesses = new Set<Promise<void>>();
|
|
const activeBatches = new Set<string>();
|
|
const DEFAULT_BATCH_WINDOW_SECONDS: number = 24 * 60 * 60;
|
|
const BATCH_RETRY_DURATION_MS: number =
|
|
Number.parseInt(process.env.BATCH_RETRY_DURATION_MS ?? "", 10) || 24 * 60 * 60 * 1_000;
|
|
const BATCH_BACKOFF_BASE_MS: number =
|
|
Number.parseInt(process.env.BATCH_BACKOFF_BASE_MS ?? "", 10) || 5_000;
|
|
const BATCH_BACKOFF_MAX_MS: number =
|
|
Number.parseInt(process.env.BATCH_BACKOFF_MAX_MS ?? "", 10) || 3_600_000;
|
|
const BATCH_MAX_CONCURRENT: number =
|
|
Number.parseInt(process.env.BATCH_MAX_CONCURRENT ?? "", 10) || 1;
|
|
|
|
interface BatchRequestItem {
|
|
body: Record<string, unknown>;
|
|
customId: string | null;
|
|
lineNumber: number;
|
|
method: "POST";
|
|
url: SupportedBatchEndpoint;
|
|
}
|
|
|
|
export function initBatchProcessor() {
|
|
if (pollInterval) return pollInterval;
|
|
console.log("[BATCH] Initializing batch processor polling...");
|
|
|
|
pollInterval = setInterval(async (): Promise<void> => {
|
|
if (isProcessing) return;
|
|
try {
|
|
isProcessing = true;
|
|
await processPendingBatches();
|
|
} catch (err) {
|
|
console.error("[BATCH] Polling error:", err);
|
|
} finally {
|
|
isProcessing = false;
|
|
}
|
|
}, 10000); // Poll every 10s
|
|
return pollInterval;
|
|
}
|
|
|
|
export function stopBatchProcessor(): void {
|
|
if (pollInterval) {
|
|
clearInterval(pollInterval);
|
|
pollInterval = null;
|
|
console.log("[BATCH] Stopped batch processor polling.");
|
|
}
|
|
}
|
|
|
|
export async function processPendingBatches(): Promise<void> {
|
|
const pending = getPendingBatches();
|
|
|
|
// Phase 1: Stale recovery — in_progress/finalizing batches not in activeBatches
|
|
// are from a previous session; reset checkpointed batches to validating so they
|
|
// can be completed without replaying already-dispatched items.
|
|
for (const batch of pending) {
|
|
if (batch.status === "in_progress" || batch.status === "finalizing") {
|
|
if (!activeBatches.has(batch.id)) {
|
|
recoverStaleBatch(batch);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Phase 2: Process actions respecting concurrency limit
|
|
const remaining = getPendingBatches(); // re-fetch after recovery updates
|
|
let activeCount = activeBatches.size;
|
|
|
|
for (const batch of remaining) {
|
|
if (batch.status === "validating") {
|
|
if (activeCount >= BATCH_MAX_CONCURRENT) {
|
|
console.log(
|
|
`[BATCH] Concurrency limit ${BATCH_MAX_CONCURRENT} reached, deferring batch ${batch.id}`
|
|
);
|
|
continue;
|
|
}
|
|
activeCount++;
|
|
await startBatch(batch);
|
|
} else if (batch.status === "cancelling") {
|
|
await cancelBatch(batch);
|
|
}
|
|
}
|
|
|
|
// Cleanup task: delete files for batches completed more than completionWindow ago
|
|
await cleanupExpiredBatches();
|
|
}
|
|
|
|
function recoverStaleBatch(batch: BatchRecord): void {
|
|
const checkpointCount = countBatchItemCheckpoints(batch.id);
|
|
const hasPotentialExternalEffects =
|
|
batch.requestCountsTotal > 0 ||
|
|
batch.requestCountsCompleted > 0 ||
|
|
batch.requestCountsFailed > 0 ||
|
|
batch.status === "finalizing";
|
|
|
|
if (checkpointCount === 0 && hasPotentialExternalEffects) {
|
|
console.warn(
|
|
`[BATCH] Stale batch ${batch.id} has no item checkpoints; failing instead of replaying provider calls`
|
|
);
|
|
failBatch(
|
|
batch.id,
|
|
"Cannot safely recover stale batch because item checkpoints are unavailable; create a new batch to retry intentionally."
|
|
);
|
|
return;
|
|
}
|
|
|
|
console.log(`[BATCH] Recovering stale batch ${batch.id} (${batch.status}) → validating`);
|
|
|
|
if (batch.outputFileId) {
|
|
deleteFile(batch.outputFileId);
|
|
}
|
|
if (batch.errorFileId) {
|
|
deleteFile(batch.errorFileId);
|
|
}
|
|
|
|
updateBatch(batch.id, {
|
|
status: "validating",
|
|
inProgressAt: null,
|
|
finalizingAt: null,
|
|
outputFileId: null,
|
|
errorFileId: null,
|
|
...(checkpointCount === 0
|
|
? {
|
|
requestCountsCompleted: 0,
|
|
requestCountsFailed: 0,
|
|
}
|
|
: {}),
|
|
});
|
|
}
|
|
|
|
function parseBatchWindowSeconds(window: string | null | undefined): number {
|
|
if (!window) return DEFAULT_BATCH_WINDOW_SECONDS;
|
|
const match = /^(\d+)([hdm])$/.exec(window);
|
|
if (!match) return DEFAULT_BATCH_WINDOW_SECONDS;
|
|
|
|
const value = Number.parseInt(match[1], 10);
|
|
const unit = match[2];
|
|
if (unit === "h") return value * 3600;
|
|
if (unit === "d") return value * 86400;
|
|
if (unit === "m") return value * 60;
|
|
return DEFAULT_BATCH_WINDOW_SECONDS;
|
|
}
|
|
|
|
function getBatchOutputExpiresAt(batch: BatchRecord): number | null {
|
|
if (
|
|
batch.outputExpiresAfterAnchor === "created_at" &&
|
|
typeof batch.outputExpiresAfterSeconds === "number" &&
|
|
batch.outputExpiresAfterSeconds > 0
|
|
) {
|
|
return batch.createdAt + batch.outputExpiresAfterSeconds;
|
|
}
|
|
|
|
const completionTime: number =
|
|
batch.completedAt || batch.failedAt || batch.cancelledAt || batch.expiredAt;
|
|
if (!completionTime) return null;
|
|
// Default: batch output files expire 30 days after completion
|
|
return completionTime + DEFAULT_BATCH_EXPIRATION_SECONDS;
|
|
}
|
|
|
|
function resolveBatchApiKeyValue(batch: Pick<BatchRecord, "apiKeyId">, apiKeyRow: any): any {
|
|
if (typeof apiKeyRow?.key === "string" && apiKeyRow.key.length > 0) {
|
|
return apiKeyRow.key;
|
|
}
|
|
if (batch.apiKeyId === "env-key") {
|
|
return process.env.OMNIROUTE_API_KEY || process.env.ROUTER_API_KEY || null;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
export function parseBatchItems(
|
|
content: Buffer,
|
|
batchEndpoint: SupportedBatchEndpoint
|
|
): { items: BatchRequestItem[]; error: null } | { items: null; error: string } {
|
|
const lines = content
|
|
.toString()
|
|
.split("\n")
|
|
.map((line) => line.trim())
|
|
.filter(Boolean);
|
|
|
|
const items: BatchRequestItem[] = [];
|
|
for (const [index, line] of lines.entries()) {
|
|
let parsed: Record<string, unknown>;
|
|
try {
|
|
parsed = JSON.parse(line);
|
|
} catch {
|
|
return { items: null, error: `Line ${index + 1} is not valid JSON` };
|
|
}
|
|
|
|
const method = String(parsed.method || "POST").toUpperCase();
|
|
const url = parsed.url;
|
|
const body = parsed.body;
|
|
|
|
if (method !== "POST") {
|
|
return {
|
|
items: null,
|
|
error: `Line ${index + 1} uses unsupported method ${method}; only POST is supported`,
|
|
};
|
|
}
|
|
if (url !== batchEndpoint) {
|
|
return {
|
|
items: null,
|
|
error: `Line ${index + 1} url ${String(url)} does not match batch endpoint ${batchEndpoint}`,
|
|
};
|
|
}
|
|
if (!body || typeof body !== "object" || Array.isArray(body)) {
|
|
return { items: null, error: `Line ${index + 1} must include a JSON object body` };
|
|
}
|
|
|
|
items.push({
|
|
body: body as Record<string, unknown>,
|
|
customId: typeof parsed.custom_id === "string" ? parsed.custom_id : null,
|
|
lineNumber: index + 1,
|
|
method: "POST",
|
|
url: batchEndpoint,
|
|
});
|
|
}
|
|
|
|
return { items, error: null };
|
|
}
|
|
|
|
async function cleanupExpiredBatches(): Promise<void> {
|
|
try {
|
|
const now = Math.floor(Date.now() / 1_000);
|
|
const batches = getTerminalBatches();
|
|
|
|
// Delete files for terminal batches that have exceeded their completion window
|
|
for (const batch of batches) {
|
|
const completionTime =
|
|
batch.completedAt || batch.failedAt || batch.cancelledAt || batch.expiredAt;
|
|
// Input files expire 30 days after batch completion
|
|
const inputExpiresAt =
|
|
completionTime && batch.inputFileId
|
|
? completionTime + DEFAULT_BATCH_EXPIRATION_SECONDS
|
|
: null;
|
|
const outputExpiresAt = getBatchOutputExpiresAt(batch);
|
|
|
|
if (batch.inputFileId && inputExpiresAt && now > inputExpiresAt) {
|
|
deleteFile(batch.inputFileId);
|
|
}
|
|
if (batch.outputFileId && outputExpiresAt && now > outputExpiresAt) {
|
|
deleteFile(batch.outputFileId);
|
|
}
|
|
if (batch.errorFileId && outputExpiresAt && now > outputExpiresAt) {
|
|
deleteFile(batch.errorFileId);
|
|
}
|
|
}
|
|
|
|
// Expire validating batches that have exceeded their completion window
|
|
for (const batch of getPendingBatches()) {
|
|
if (batch.status === "validating") {
|
|
const windowSeconds = parseBatchWindowSeconds(batch.completionWindow);
|
|
if (now - batch.createdAt > windowSeconds) {
|
|
updateBatch(batch.id, { status: "expired", expiredAt: now });
|
|
}
|
|
}
|
|
}
|
|
|
|
// Cleanup orphan files (batch-purpose files stuck in validating after 48h)
|
|
// Use asc order so oldest files are processed first; use a high limit to avoid missing old orphans.
|
|
const allFiles = listFiles({ order: "asc", limit: 100 });
|
|
for (const file of allFiles) {
|
|
if (file.purpose === "batch" && now - file.createdAt > DEFAULT_BATCH_EXPIRATION_SECONDS) {
|
|
deleteFile(file.id);
|
|
}
|
|
}
|
|
} catch (err) {
|
|
console.error("[BATCH] Cleanup error:", err);
|
|
}
|
|
}
|
|
|
|
async function startBatch(batch: any): Promise<void> {
|
|
console.log(`[BATCH] Starting batch ${batch.id}`);
|
|
|
|
const content = getFileContent(batch.inputFileId);
|
|
if (!content) {
|
|
failBatch(batch.id, "Input file content not found");
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const parsedItems = parseBatchItems(content, batch.endpoint);
|
|
if (parsedItems.error) {
|
|
// Set total count even on validation failure so UI shows correct numbers
|
|
const lines = content
|
|
.toString()
|
|
.split("\n")
|
|
.map((line) => line.trim())
|
|
.filter(Boolean);
|
|
updateBatch(batch.id, {
|
|
requestCountsTotal: lines.length,
|
|
requestCountsFailed: lines.length, // All failed due to validation error
|
|
});
|
|
failBatch(batch.id, parsedItems.error);
|
|
return;
|
|
}
|
|
const total = parsedItems.items.length;
|
|
|
|
console.log(`[BATCH] Batch ${batch.id} contains (${total} items)`);
|
|
|
|
ensureBatchItemCheckpoints(batch.id, parsedItems.items);
|
|
|
|
updateBatch(batch.id, {
|
|
status: "in_progress",
|
|
inProgressAt: Math.floor(Date.now() / 1000),
|
|
requestCountsTotal: total,
|
|
});
|
|
|
|
activeBatches.add(batch.id);
|
|
|
|
// Fire-and-forget: process items in the background so the poll loop isn't blocked.
|
|
// isProcessing prevents a second poll tick from overlapping.
|
|
const p = processBatchItems(batch, parsedItems.items).catch((err) => {
|
|
console.error(`[BATCH] Critical error in processBatchItems for ${batch.id}:`, err);
|
|
failBatch(batch.id, String(err));
|
|
});
|
|
activeProcesses.add(p);
|
|
p.finally(() => {
|
|
activeProcesses.delete(p);
|
|
activeBatches.delete(batch.id);
|
|
});
|
|
} catch (err) {
|
|
console.error(`[BATCH] Error starting batch ${batch.id}:`, err);
|
|
failBatch(batch.id, err instanceof Error ? err.message : String(err));
|
|
}
|
|
}
|
|
|
|
let prevHeaders: Headers | null = null;
|
|
let prevHeadersTimestamp: number = 0;
|
|
const HEADERS_CACHE_TTL_MS = 60_000;
|
|
|
|
async function processBatchItems(batch: BatchRecord, items: BatchRequestItem[]): Promise<void> {
|
|
const state = createBatchState(batch);
|
|
const checkpoints = new Map<number, BatchItemCheckpoint>(
|
|
listBatchItemCheckpoints(batch.id).map((checkpoint) => [checkpoint.lineNumber, checkpoint])
|
|
);
|
|
|
|
const apiKey = await resolveApiKey(batch);
|
|
|
|
for (const item of items) {
|
|
if (isBatchCancelled(batch.id)) break;
|
|
|
|
const checkpoint = checkpoints.get(item.lineNumber);
|
|
if (checkpoint && applyRecoveredCheckpoint(batch.id, item, checkpoint, state)) {
|
|
maybePersistProgress(batch.id, state);
|
|
continue;
|
|
}
|
|
|
|
const cachedHeaders =
|
|
prevHeaders && Date.now() - prevHeadersTimestamp < HEADERS_CACHE_TTL_MS ? prevHeaders : null;
|
|
if (cachedHeaders) {
|
|
const delay = maybeThrottle(cachedHeaders);
|
|
if (delay) {
|
|
await sleep(delay);
|
|
}
|
|
}
|
|
|
|
markBatchItemProcessing(batch.id, item);
|
|
|
|
try {
|
|
const response = await processSingleItemWithRetry(item, apiKey);
|
|
let responseBody: unknown;
|
|
try {
|
|
responseBody = await response.clone().json();
|
|
} catch {
|
|
responseBody = await response.text();
|
|
}
|
|
|
|
// Record the item's result so finalization can emit output/error files.
|
|
// The output file format expects entries like:
|
|
// { id, custom_id, response: { status_code, body } }
|
|
const wrapped = {
|
|
id: `req_${uuidv4().replaceAll("-", "")}`,
|
|
custom_id: item.customId ?? null,
|
|
response: {
|
|
status_code: response.status,
|
|
body: responseBody,
|
|
},
|
|
};
|
|
|
|
markBatchItemResult(batch.id, item, wrapped);
|
|
state.results.push(wrapped);
|
|
applyItemResult(state, response.status, responseBody);
|
|
prevHeaders = response.headers;
|
|
prevHeadersTimestamp = Date.now();
|
|
} catch (exception) {
|
|
// Track processing-level errors separately (items that failed to be processed)
|
|
const error = { custom_id: item.customId ?? null, error: String(exception) };
|
|
markBatchItemError(batch.id, item, error);
|
|
state.errors.push(error);
|
|
state.failed++;
|
|
prevHeaders = null;
|
|
prevHeadersTimestamp = 0;
|
|
}
|
|
|
|
maybePersistProgress(batch.id, state);
|
|
}
|
|
|
|
return finalizeBatch(batch.id, state.results, state.errors);
|
|
}
|
|
|
|
function applyRecoveredCheckpoint(
|
|
batchId: string,
|
|
item: BatchRequestItem,
|
|
checkpoint: BatchItemCheckpoint,
|
|
state: ReturnType<typeof createBatchState>
|
|
): boolean {
|
|
if (checkpoint.status === "completed" && checkpoint.result) {
|
|
state.results.push(checkpoint.result);
|
|
applyItemResult(
|
|
state,
|
|
checkpoint.result.response?.status_code ?? 500,
|
|
checkpoint.result.response?.body
|
|
);
|
|
return true;
|
|
}
|
|
|
|
if (checkpoint.status === "errored" && checkpoint.error) {
|
|
state.errors.push(checkpoint.error);
|
|
state.failed++;
|
|
return true;
|
|
}
|
|
|
|
if (checkpoint.status === "processing") {
|
|
const error = {
|
|
custom_id: item.customId ?? null,
|
|
error:
|
|
"Batch item was interrupted before its provider response was recorded; it was not replayed to avoid duplicate provider work.",
|
|
};
|
|
markBatchItemError(batchId, item, error);
|
|
state.errors.push(error);
|
|
state.failed++;
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
function isBatchCancelled(batchId: string): boolean {
|
|
const current = getBatch(batchId);
|
|
|
|
return !current || current.status === "cancelling" || current.status === "cancelled";
|
|
}
|
|
|
|
async function resolveApiKey(batch: BatchRecord): Promise<any> {
|
|
const apiKeyRow = batch.apiKeyId ? await getApiKeyById(batch.apiKeyId) : null;
|
|
return resolveBatchApiKeyValue(batch, apiKeyRow);
|
|
}
|
|
|
|
async function processSingleItemWithRetry(item: BatchRequestItem, apiKey: string) {
|
|
// Time-based retry limit: individual batch items can retry for up to 24 hours.
|
|
// This accommodates large batches against heavily rate-limited providers.
|
|
const MAX_RETRY_DURATION_MS = BATCH_RETRY_DURATION_MS;
|
|
const maxRetries = 200; // safety ceiling — time limit should kick in first
|
|
const retryStartedAt = Date.now();
|
|
|
|
let response: Response = null;
|
|
for (let attempt = 1; attempt <= maxRetries; attempt++) {
|
|
// If the previous attempt got a response, check for rate limit headers and throttle if needed before retrying.
|
|
if (response) {
|
|
const delay = maybeThrottle(response.headers);
|
|
if (delay) {
|
|
await sleep(delay);
|
|
}
|
|
}
|
|
|
|
response = await processSingleItem(item, apiKey);
|
|
|
|
if (
|
|
(response.status === 429 || response.status === 502 || response.status === 504) &&
|
|
attempt < maxRetries
|
|
) {
|
|
// Bail if we've been retrying for longer than the allowed window
|
|
if (Date.now() - retryStartedAt >= MAX_RETRY_DURATION_MS) {
|
|
console.warn(
|
|
`[BATCH] Item ${item.customId ?? "(no id)"} exceeded 24h retry window after ${attempt} attempts — giving up`
|
|
);
|
|
return response;
|
|
}
|
|
const delay = getRetryDelayMs(response.headers) ?? getBackoffDelayMs(attempt);
|
|
await sleep(delay);
|
|
continue;
|
|
}
|
|
|
|
return response;
|
|
}
|
|
}
|
|
|
|
async function processSingleItem(item: BatchRequestItem, apiKey: string) {
|
|
const body = buildRequestBody(item);
|
|
|
|
return await dispatch.dispatchBatchApiRequest({
|
|
endpoint: item.url,
|
|
body,
|
|
apiKey,
|
|
});
|
|
}
|
|
|
|
export function buildRequestBody(item: BatchRequestItem) {
|
|
const isChatEndpoint = ![
|
|
"/v1/embeddings",
|
|
"/v1/moderations",
|
|
"/v1/images/generations",
|
|
"/v1/images/edits",
|
|
"/v1/videos",
|
|
"/v1/videos/generations",
|
|
].includes(item.url);
|
|
|
|
return {
|
|
...item.body,
|
|
...(isChatEndpoint ? { stream: false } : {}),
|
|
};
|
|
}
|
|
|
|
function getBackoffDelayMs(attempt: number) {
|
|
const baseMs = BATCH_BACKOFF_BASE_MS;
|
|
const maxMs = BATCH_BACKOFF_MAX_MS;
|
|
|
|
// exponential: 2^attempt * base
|
|
const exp = Math.min(maxMs, baseMs * 2 ** attempt);
|
|
|
|
// jitter ±20%
|
|
const jitterFactor = 1 + (Math.random() * 0.4 - 0.2);
|
|
|
|
return Math.floor(exp * jitterFactor);
|
|
}
|
|
|
|
function sleep(ms: number) {
|
|
return new Promise((res) => setTimeout(res, ms));
|
|
}
|
|
function getRetryDelayMs(headers: Headers): number | null {
|
|
const retryAfter = headers.get("retry-after");
|
|
if (retryAfter) {
|
|
const seconds = Number(retryAfter);
|
|
if (!Number.isNaN(seconds)) {
|
|
return seconds * 1_000;
|
|
}
|
|
|
|
// fallback: HTTP-date
|
|
const date = new Date(retryAfter).getTime();
|
|
if (!Number.isNaN(date)) {
|
|
return Math.max(0, date - Date.now());
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
export function maybeThrottle(headers: Headers): number | null {
|
|
// Mistral reports these headers from their API
|
|
const remainingReq = toNumber(headers.get("x-ratelimit-remaining-req-minute"));
|
|
const limitReq = toNumber(headers.get("x-ratelimit-limit-req-minute"));
|
|
|
|
const remainingTokens = toNumber(headers.get("x-ratelimit-remaining-tokens-minute"));
|
|
const cost = toNumber(headers.get("x-ratelimit-tokens-query-cost"));
|
|
|
|
let pressures: number[] = [];
|
|
|
|
// Request pressureRemaining
|
|
if (remainingReq !== null && limitReq !== null) {
|
|
if (limitReq > 0) {
|
|
pressures.push(remainingReq / limitReq);
|
|
}
|
|
}
|
|
|
|
// Token pressureRemaining
|
|
if (remainingTokens !== null && cost !== null) {
|
|
if (remainingTokens + cost > 0) {
|
|
pressures.push(remainingTokens / (remainingTokens + cost));
|
|
}
|
|
}
|
|
|
|
if (pressures.length === 0) {
|
|
console.log("[BATCH] Throttle check - no rate-limit headers present");
|
|
return null;
|
|
} else {
|
|
const tokenTotal = remainingTokens != null && cost != null ? remainingTokens + cost : null;
|
|
console.log(
|
|
`[BATCH] Throttle check - Request pressure: ${remainingReq ?? "n/a"}/${limitReq ?? "n/a"}, Token pressure: ${remainingTokens ?? "n/a"}/${tokenTotal ?? "n/a"}`
|
|
);
|
|
}
|
|
|
|
const pressureRemaining = Math.min(...pressures);
|
|
|
|
const delay = throttleDelay(pressureRemaining);
|
|
if (delay !== null) {
|
|
console.log(
|
|
`[BATCH] Throttling next request with delay of ${Math.round(delay)}ms (pressure remaining: ${(pressureRemaining * 100).toFixed(2)}%)`
|
|
);
|
|
}
|
|
return delay;
|
|
}
|
|
|
|
function throttleDelay(pressure: number): number | null {
|
|
if (pressure >= 0.2) return null;
|
|
|
|
const severity = (0.2 - pressure) / 0.2;
|
|
|
|
const delay = Math.pow(severity, 2) * 30_000;
|
|
|
|
return 200 + delay + Math.random() * 1000;
|
|
}
|
|
|
|
const toNumber = (v: string | null) => {
|
|
if (v === null || v.trim() === "") return null;
|
|
const n = Number(v);
|
|
return Number.isFinite(n) ? n : null;
|
|
};
|
|
|
|
function createBatchState(batch: BatchRecord) {
|
|
return {
|
|
results: [],
|
|
errors: [],
|
|
completed: 0,
|
|
failed: 0,
|
|
tokens: {
|
|
input: 0,
|
|
output: 0,
|
|
reasoning: 0,
|
|
},
|
|
model: batch.model || null,
|
|
};
|
|
}
|
|
|
|
function applyItemResult(state: any, statusCode: number, body: any): void {
|
|
if (statusCode >= 400 || body?.error) {
|
|
state.failed++;
|
|
} else {
|
|
state.completed++;
|
|
|
|
if (body?.usage) {
|
|
state.tokens.input += body.usage.prompt_tokens || body.usage.input_tokens || 0;
|
|
state.tokens.output += body.usage.completion_tokens || body.usage.output_tokens || 0;
|
|
state.tokens.reasoning += body.usage.completion_tokens_details?.reasoning_tokens || 0;
|
|
}
|
|
|
|
if (!state.model && body?.model) {
|
|
state.model = body.model;
|
|
}
|
|
}
|
|
}
|
|
|
|
function maybePersistProgress(batchId: string, state: any): void {
|
|
// Persist basic progress (completed/failed counts + model) on every item so
|
|
// the UI can show up-to-date progress even for small batches.
|
|
try {
|
|
updateBatch(batchId, {
|
|
requestCountsCompleted: state.completed,
|
|
requestCountsFailed: state.failed,
|
|
model: state.model,
|
|
});
|
|
} catch (err) {
|
|
console.error(`[BATCH] Failed to persist progress for ${batchId}:`, err);
|
|
}
|
|
|
|
// Persist richer usage/model info less frequently to avoid excessive writes.
|
|
const total = state.completed + state.failed;
|
|
if (total % 50 !== 0) return;
|
|
|
|
try {
|
|
updateBatch(batchId, {
|
|
requestCountsCompleted: state.completed,
|
|
requestCountsFailed: state.failed,
|
|
model: state.model,
|
|
usage: {
|
|
input_tokens: state.tokens.input,
|
|
output_tokens: state.tokens.output,
|
|
total_tokens: state.tokens.input + state.tokens.output,
|
|
input_tokens_details: { cached_tokens: 0 },
|
|
output_tokens_details: { reasoning_tokens: state.tokens.reasoning },
|
|
},
|
|
});
|
|
} catch (err) {
|
|
console.error(`[BATCH] Failed to persist extended progress for ${batchId}:`, err);
|
|
}
|
|
}
|
|
|
|
async function finalizeBatch(
|
|
batchId: string,
|
|
results: any[],
|
|
itemsWithErrors: any[]
|
|
): Promise<void> {
|
|
const current = getBatch(batchId);
|
|
|
|
if (handleCancellation(batchId, current)) return;
|
|
|
|
// Mark as finalizing first
|
|
markFinalizing(batchId);
|
|
|
|
// Compute counts from results
|
|
const successes = results.filter(
|
|
(r) =>
|
|
typeof r.response?.status_code === "number" &&
|
|
r.response.status_code < 400 &&
|
|
!r.response.body?.error
|
|
);
|
|
const failuresFromResults = results.filter(
|
|
(r) =>
|
|
(typeof r.response?.status_code === "number" && r.response.status_code >= 400) ||
|
|
r.response?.body?.error
|
|
);
|
|
const processingErrors = itemsWithErrors || [];
|
|
|
|
const completedCount = successes.length;
|
|
const failedCount = failuresFromResults.length + processingErrors.length;
|
|
const totalCount = current?.requestCountsTotal || completedCount + failedCount;
|
|
|
|
// Aggregate usage from per-item responses when available
|
|
let inputTokens = 0;
|
|
let outputTokens = 0;
|
|
let reasoningTokens = 0;
|
|
|
|
for (const r of results) {
|
|
try {
|
|
const body = r.response?.body;
|
|
if (!body) continue;
|
|
const usage = body.usage || {};
|
|
inputTokens += usage.prompt_tokens ?? usage.input_tokens ?? usage.total_tokens ?? 0;
|
|
outputTokens += usage.completion_tokens ?? 0;
|
|
reasoningTokens += usage.completion_tokens_details?.reasoning_tokens ?? 0;
|
|
} catch (err) {
|
|
console.error("Failed to aggregate usage for batch", batchId, err);
|
|
}
|
|
}
|
|
|
|
const model =
|
|
results.find((r) => r.response?.body?.model)?.response?.body?.model || current?.model || null;
|
|
|
|
const completionTime = now();
|
|
|
|
// Persist final counts and (approximate) usage so UI shows correct numbers
|
|
try {
|
|
updateBatch(batchId, {
|
|
requestCountsTotal: totalCount,
|
|
requestCountsCompleted: completedCount,
|
|
requestCountsFailed: failedCount,
|
|
model,
|
|
usage: {
|
|
input_tokens: inputTokens,
|
|
output_tokens: outputTokens,
|
|
total_tokens: inputTokens + outputTokens,
|
|
input_tokens_details: { cached_tokens: 0 },
|
|
output_tokens_details: { reasoning_tokens: reasoningTokens },
|
|
},
|
|
// also set completedAt so file expiration calculation can use it
|
|
completedAt: completionTime,
|
|
});
|
|
} catch (err) {
|
|
console.error(`[BATCH] Failed to persist final progress for ${batchId}:`, err);
|
|
}
|
|
|
|
// Re-read the batch (with completedAt set) so file creation sees a completion timestamp
|
|
const batchForFiles = getBatch(batchId);
|
|
|
|
const outputFileId = createSuccessFile(batchId, batchForFiles, results);
|
|
const errorFileId = createErrorFile(batchId, batchForFiles, results, itemsWithErrors);
|
|
|
|
completeBatch(batchId, outputFileId, errorFileId);
|
|
}
|
|
|
|
function handleCancellation(batchId: string, current: any): boolean {
|
|
if (!current) return true;
|
|
|
|
if (current.status === "cancelling") {
|
|
updateBatch(batchId, {
|
|
status: "cancelled",
|
|
cancelledAt: now(),
|
|
});
|
|
return true;
|
|
}
|
|
|
|
return current.status === "cancelled";
|
|
}
|
|
|
|
function markFinalizing(batchId: string): void {
|
|
updateBatch(batchId, {
|
|
status: "finalizing",
|
|
finalizingAt: now(),
|
|
});
|
|
}
|
|
|
|
function completeBatch(
|
|
batchId: string,
|
|
outputFileId: string | null,
|
|
errorFileId: string | null
|
|
): void {
|
|
updateBatch(batchId, {
|
|
status: "completed",
|
|
completedAt: now(),
|
|
outputFileId,
|
|
errorFileId,
|
|
});
|
|
|
|
const b = getBatch(batchId);
|
|
const total = b?.requestCountsTotal ?? "?";
|
|
console.log(`[BATCH] Completed batch ${batchId} (${total} items)`);
|
|
}
|
|
|
|
function now(): number {
|
|
return Math.floor(Date.now() / 1_000);
|
|
}
|
|
|
|
function createSuccessFile(batchId: string, current: any, results: any[]): string | null {
|
|
const successes = results.filter((r) => r.response.status_code < 400 && !r.response.body?.error);
|
|
|
|
if (successes.length === 0) return null;
|
|
|
|
const content = toJsonl(successes);
|
|
|
|
const file = createFile({
|
|
bytes: Buffer.byteLength(content),
|
|
filename: `batch_${batchId}_output.jsonl`,
|
|
purpose: "batch_output",
|
|
content: Buffer.from(content),
|
|
apiKeyId: current?.apiKeyId,
|
|
expiresAt: getBatchOutputExpiresAt(current),
|
|
});
|
|
|
|
return file.id;
|
|
}
|
|
|
|
function createErrorFile(
|
|
batchId: string,
|
|
current: any,
|
|
results: any[],
|
|
itemsWithErrors: any[]
|
|
): string | null {
|
|
const failures = results.filter((r) => r.response.status_code >= 400 || r.response.body?.error);
|
|
|
|
const processErrors = itemsWithErrors.map((e) => ({
|
|
id: `batch_req_${uuidv4().replaceAll("-", "")}`,
|
|
custom_id: e.custom_id,
|
|
response: null,
|
|
error: { message: e.error, type: "batch_process_error" },
|
|
}));
|
|
|
|
const allFailures = [...failures, ...processErrors];
|
|
|
|
if (allFailures.length === 0) return null;
|
|
|
|
const content = toJsonl(allFailures);
|
|
|
|
const file = createFile({
|
|
bytes: Buffer.byteLength(content),
|
|
filename: `batch_${batchId}_error.jsonl`,
|
|
purpose: "batch_output",
|
|
content: Buffer.from(content),
|
|
apiKeyId: current?.apiKeyId,
|
|
expiresAt: getBatchOutputExpiresAt(current),
|
|
});
|
|
|
|
return file.id;
|
|
}
|
|
|
|
function toJsonl(items: any[]): string {
|
|
return items.map((i) => JSON.stringify(i)).join("\n");
|
|
}
|
|
|
|
async function cancelBatch(batch: any): Promise<void> {
|
|
updateBatch(batch.id, {
|
|
status: "cancelled",
|
|
cancelledAt: Math.floor(Date.now() / 1000),
|
|
});
|
|
console.log(`[BATCH] Cancelled batch ${batch.id}`);
|
|
}
|
|
|
|
function failBatch(batchId: string, reason: string): void {
|
|
updateBatch(batchId, {
|
|
status: "failed",
|
|
failedAt: Math.floor(Date.now() / 1000),
|
|
errors: [{ message: reason }],
|
|
});
|
|
activeBatches.delete(batchId);
|
|
}
|
|
|
|
export async function waitForAllBatches(): Promise<void> {
|
|
await Promise.all(Array.from(activeProcesses));
|
|
}
|
|
|
|
// Test helpers
|
|
export function getCachedHeaders(): { headers: Headers | null; timestamp: number } {
|
|
return { headers: prevHeaders, timestamp: prevHeadersTimestamp };
|
|
}
|
|
export function resetCachedHeaders(): void {
|
|
prevHeaders = null;
|
|
prevHeadersTimestamp = 0;
|
|
}
|
|
export function resetBatchProcessorState(): void {
|
|
activeBatches.clear();
|
|
activeProcesses.clear();
|
|
isProcessing = false;
|
|
prevHeaders = null;
|
|
prevHeadersTimestamp = 0;
|
|
}
|