haixunMaster/lib/search/quota.ts

44 lines
1.2 KiB
TypeScript

import fs from "fs";
import path from "path";
const QUOTA_DIR = path.join(process.cwd(), "data", "search-quota");
function todayKey(): string {
return new Date().toISOString().slice(0, 10);
}
function quotaFile(provider: string): string {
fs.mkdirSync(QUOTA_DIR, { recursive: true });
return path.join(QUOTA_DIR, `${provider}-${todayKey()}.json`);
}
interface QuotaState {
count: number;
}
export function getDailyQuotaUsed(provider: string): number {
try {
const file = quotaFile(provider);
if (!fs.existsSync(file)) return 0;
const state = JSON.parse(fs.readFileSync(file, "utf8")) as QuotaState;
return state.count ?? 0;
} catch {
return 0;
}
}
export function incrementDailyQuota(provider: string): number {
const used = getDailyQuotaUsed(provider) + 1;
fs.writeFileSync(quotaFile(provider), JSON.stringify({ count: used } satisfies QuotaState));
return used;
}
export function canUseDailyQuota(provider: string, limit: number): boolean {
return getDailyQuotaUsed(provider) < limit;
}
export function resetQuotaForTests(): void {
if (fs.existsSync(QUOTA_DIR)) {
fs.rmSync(QUOTA_DIR, { recursive: true, force: true });
}
}