ratelimit.js
2.07 KB
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
/**
* `createAiLimit` —— 给烧钱的 AI 端点（LLM 续写 / TTS 合成 / 生图）挂 per-IP / 匿名基础
* 限流的 Fastify preHandler 工厂，超限返结构化 `429`。
*
* 只做「按端点类别构造 RateQuery → consume → 超限成形 429」，桶数学 / store 在 @drama/ratelimit
* 内。subject 取匿名 device-id（`x-device-id`）优先，否则回退客户端 IP（per-IP 防脚本薅穿）。
*/
import { parse } from 'valibot'
import { ApiError } from '@drama/contracts/platform'
/**
* @param {{ limiter: { consume: Function } }} deps @drama/ratelimit 的 limiter 实例
* @returns {(endpoint: 'llm'|'tts'|'image') => import('fastify').preHandlerHookHandler}
*/
export function createAiLimit({ limiter }) {
if (!limiter || typeof limiter.consume !== 'function') {
throw new Error('createAiLimit: 缺少注入的 limiter（需实现 consume）')
}
return function aiLimit(endpoint) {
return async function preHandler(request, reply) {
const subject = subjectOf(request)
const decision = await limiter.consume(request.ctx, { endpoint, subject })
if (!decision.allowed) {
const retryAfterSec = Math.ceil(decision.retryAfterMs / 1000)
reply
.code(429)
.header('retry-after', String(retryAfterSec))
.header('x-ratelimit-remaining', String(Math.max(0, decision.remaining)))
return reply.send(
parse(ApiError, {
error: {
code: 'rate_limited',
message: `Too many ${endpoint} requests, retry after ${retryAfterSec}s`,
retryable: true,
details: { endpoint, retryAfterMs: decision.retryAfterMs, ruleId: decision.ruleId },
},
}),
)
}
reply.header('x-ratelimit-remaining', String(Math.max(0, decision.remaining)))
}
}
}
/**
* @param {import('fastify').FastifyRequest} request
* @returns {string}
*/
function subjectOf(request) {
const dev = request.headers['x-device-id']
if (typeof dev === 'string' && dev) return `anon:${dev}`
return `ip:${request.ip || 'unknown'}`
}