File size: 3,238 Bytes
20f83d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import { getKeyPrefix } from './redis';
import { PRO_DAILY_QUOTA_TTL_SECONDS, secondsUntilUtcMidnight } from './pro-mcp-token';
export const DIRECT_LLM_DAILY_QUOTA_LIMIT = 50;
export const DIRECT_LLM_REDIS_UNAVAILABLE_RETRY_AFTER_SECONDS = 30;
export const DIRECT_LLM_GATEWAY_QUOTA_PATHS = new Set<string>([
'/api/intelligence/v1/classify-event',
'/api/intelligence/v1/deduct-situation',
'/api/intelligence/v1/get-country-intel-brief',
'/api/market/v1/analyze-stock',
'/api/news/v1/summarize-article',
]);
export const DIRECT_LLM_SELF_METERED_QUOTA_PATHS = new Set<string>([
'/api/chat-analyst',
]);
export const DIRECT_LLM_QUOTA_PATHS = new Set<string>([
...DIRECT_LLM_GATEWAY_QUOTA_PATHS,
...DIRECT_LLM_SELF_METERED_QUOTA_PATHS,
]);
export type DirectLlmQuotaReservation =
| { ok: true; newCount: number; rollback: () => Promise<void> }
| {
ok: false;
reason: 'cap-exceeded' | 'redis-unavailable';
floor?: number;
retryAfterSec: number;
};
export type DirectLlmQuotaPipeline = (
commands: Array<Array<string | number>>,
) => Promise<Array<{ result?: unknown }>>;
export function directLlmDailyQuotaKey(userId: string, date?: Date): string {
if (!userId) return '';
const d = date ?? new Date();
const yyyy = d.getUTCFullYear();
const mm = String(d.getUTCMonth() + 1).padStart(2, '0');
const dd = String(d.getUTCDate()).padStart(2, '0');
return `${getKeyPrefix()}llm:direct-usage:${userId}:${yyyy}-${mm}-${dd}`;
}
export async function reserveDirectLlmQuota(opts: {
userId: string;
pipeline: DirectLlmQuotaPipeline;
date?: Date;
}): Promise<DirectLlmQuotaReservation> {
const retryAfterSec = secondsUntilUtcMidnight(opts.date);
const key = directLlmDailyQuotaKey(opts.userId, opts.date);
if (!key) {
return {
ok: false,
reason: 'redis-unavailable',
retryAfterSec: DIRECT_LLM_REDIS_UNAVAILABLE_RETRY_AFTER_SECONDS,
};
}
let pipeResult: Array<{ result?: unknown }> | null;
try {
pipeResult = await opts.pipeline([
['INCR', key],
['EXPIRE', key, PRO_DAILY_QUOTA_TTL_SECONDS],
]);
} catch {
pipeResult = null;
}
if (!pipeResult || !Array.isArray(pipeResult) || pipeResult.length === 0) {
return {
ok: false,
reason: 'redis-unavailable',
retryAfterSec: DIRECT_LLM_REDIS_UNAVAILABLE_RETRY_AFTER_SECONDS,
};
}
const incrRaw = pipeResult[0]?.result;
const newCount = typeof incrRaw === 'number' ? incrRaw : Number(incrRaw);
if (!Number.isFinite(newCount) || newCount < 1) {
return {
ok: false,
reason: 'redis-unavailable',
retryAfterSec: DIRECT_LLM_REDIS_UNAVAILABLE_RETRY_AFTER_SECONDS,
};
}
let rolledBack = false;
const rollback = async (): Promise<void> => {
if (rolledBack) return;
rolledBack = true;
try {
await opts.pipeline([['DECR', key]]);
} catch {
// Best-effort: over-counting by one is the cost-protection-correct direction.
}
};
if (newCount > DIRECT_LLM_DAILY_QUOTA_LIMIT) {
await rollback();
return {
ok: false,
reason: 'cap-exceeded',
floor: DIRECT_LLM_DAILY_QUOTA_LIMIT,
retryAfterSec,
};
}
return { ok: true, newCount, rollback };
}
|