diff --git a/packages/service/common/middle/teamRateLimit.ts b/packages/service/common/middle/teamRateLimit.ts new file mode 100644 index 000000000000..453d3464bb69 --- /dev/null +++ b/packages/service/common/middle/teamRateLimit.ts @@ -0,0 +1,104 @@ +import { type ApiRequestProps } from '../../type/next'; +import { type NextApiResponse } from 'next'; +import { jsonRes } from '../response'; +import { getGlobalRedisConnection } from '../redis'; + +export interface TeamRateLimitOptions { + seconds: number; + limit: number; + keyPrefix?: string; +} + +/** + * 基于团队ID的Redis限流中间件 + * 注意:这个中间件会在请求认证前执行,所以需要从请求中提取团队ID + * @param options 限流配置 + * @returns 中间件函数 + * @example + * export default NextAPI( + * useTeamFrequencyLimit({ + * seconds: 60, + * limit: 1000, + * keyPrefix: 'chat-completions-rate-limit' + * }), + * handler + * ); + */ +export function useTeamFrequencyLimit(options: TeamRateLimitOptions) { + const { seconds, limit, keyPrefix = 'team-rate-limit' } = options; + + return async (req: ApiRequestProps, res: NextApiResponse) => { + // 尝试从请求的不同位置获取团队ID + let teamId: string | undefined; + + // 1. 从请求体中获取(最常见的情况) + if (req.body?.teamId) { + teamId = req.body.teamId; + } + // 2. 从查询参数中获取 + else if (req.query?.teamId) { + teamId = req.query.teamId as string; + } + // 3. 从Authorization header中解析(如果使用API Key) + else if (req.headers?.authorization) { + // 这里可以添加API Key解析逻辑来获取teamId + // 但为了简单起见,我们暂时跳过这种情况 + } + + if (!teamId) { + // 如果没有团队ID,跳过限流检查 + return; + } + + try { + const redis = getGlobalRedisConnection(); + const key = `${keyPrefix}:${teamId}`; + + // 使用Redis的滑动窗口限流算法 + const currentTime = Math.floor(Date.now() / 1000); + const windowStart = currentTime - seconds; + + // 使用Redis Pipeline提高性能 + const pipeline = redis.pipeline(); + + // 移除过期的请求记录 + pipeline.zremrangebyscore(key, 0, windowStart); + + // 添加当前请求记录 + pipeline.zadd(key, currentTime, `${currentTime}-${Math.random()}`); + + // 获取当前窗口内的请求数量 + pipeline.zcard(key); + + // 设置key的过期时间 + pipeline.expire(key, seconds); + + const results = await pipeline.exec(); + + if (!results) { + throw new Error('Redis pipeline execution failed'); + } + + // zcard的结果在pipeline的第三个操作中 + const currentRequestCount = results[2][1] as number; + + if (currentRequestCount > limit) { + // 超出限流,返回429错误 + const remainingTime = await redis.ttl(key); + jsonRes(res, { + code: 429, + error: `Rate limit exceeded. Maximum ${limit} requests per ${seconds} seconds for this team. Please try again in ${remainingTime} seconds.` + }); + return; + } + + // 在响应头中添加限流信息 + res.setHeader('X-RateLimit-Limit', limit); + res.setHeader('X-RateLimit-Remaining', Math.max(0, limit - currentRequestCount)); + res.setHeader('X-RateLimit-Reset', new Date(Date.now() + seconds * 1000).toISOString()); + } catch (error) { + console.error('Rate limit check failed:', error); + // Redis错误时不阻断请求,继续处理 + } + }; +} diff --git a/projects/app/src/pages/api/v1/chat/completions.ts b/projects/app/src/pages/api/v1/chat/completions.ts index 11b6cbecf37e..42a897d6f9ba 100644 --- a/projects/app/src/pages/api/v1/chat/completions.ts +++ b/projects/app/src/pages/api/v1/chat/completions.ts @@ -48,6 +48,7 @@ import { DispatchNodeResponseKeyEnum } from '@fastgpt/global/core/workflow/runti import { NextAPI } from '@/service/middleware/entry'; import { getAppLatestVersion } from '@fastgpt/service/core/app/version/controller'; import { ReadPermissionVal } from '@fastgpt/global/support/permission/constant'; +import { useTeamFrequencyLimit } from '@fastgpt/service/common/middle/teamRateLimit'; import { AppTypeEnum } from '@fastgpt/global/core/app/constants'; import { serverGetWorkflowToolRunUserQuery, @@ -449,7 +450,14 @@ async function handler(req: NextApiRequest, res: NextApiResponse) { } } } -export default NextAPI(handler); +export default NextAPI( + useTeamFrequencyLimit({ + seconds: 60, + limit: 1000, + keyPrefix: 'chat-completions-rate-limit' + }), + handler +); const authShareChat = async ({ chatId,