diff --git a/src/app/api/chat/ratelimit.test.ts b/src/app/api/chat/ratelimit.test.ts index e13a501..d2aba2b 100644 --- a/src/app/api/chat/ratelimit.test.ts +++ b/src/app/api/chat/ratelimit.test.ts @@ -24,6 +24,20 @@ describe("isRateLimited", () => { expect(result).toBe(false); }); + it("should prioritize x-real-ip over x-forwarded-for", async () => { + mockLimit.mockResolvedValue({ success: true }); + + const req = new Request("http://localhost/api/chat", { + headers: { + "x-forwarded-for": "1.1.1.1", + "x-real-ip": "2.2.2.2", + }, + }); + + await isRateLimited(req); + expect(mockLimit).toHaveBeenCalledWith("2.2.2.2"); + }); + it("should extract IP from x-forwarded-for header and call limit if x-real-ip is missing", async () => { mockLimit.mockResolvedValue({ success: true }); diff --git a/src/app/api/chat/ratelimit.ts b/src/app/api/chat/ratelimit.ts index 06f553c..6391a7d 100644 --- a/src/app/api/chat/ratelimit.ts +++ b/src/app/api/chat/ratelimit.ts @@ -7,6 +7,8 @@ export default async function isRateLimited( ): Promise { // eslint-disable-next-line @typescript-eslint/no-explicit-any let ip = "ip" in req ? (req as any).ip : null; + ip = ip ?? req.headers.get("x-real-ip"); + if (!ip) { const forwardedFor = req.headers.get("x-forwarded-for"); if (forwardedFor) { @@ -14,7 +16,8 @@ export default async function isRateLimited( ip = parts[parts.length - 1].trim(); } } - ip = ip ?? req.headers.get("x-real-ip") ?? "anonymous"; + + ip = ip ?? "anonymous"; const { success } = await ratelimit.limit(ip); return !success; } diff --git a/src/app/api/chat/request.test.ts b/src/app/api/chat/request.test.ts index 9710b52..53f15e5 100644 --- a/src/app/api/chat/request.test.ts +++ b/src/app/api/chat/request.test.ts @@ -62,7 +62,18 @@ describe("sendRequest", () => { }); await expect(sendRequest(req)).rejects.toThrow( - "Invalid request: messages must be an array.", + "Invalid request: messages must be an array and cannot be empty.", + ); + }); + + it("should throw an error if messages array is empty", async () => { + const req = new Request("http://localhost/api/chat", { + method: "POST", + body: JSON.stringify({ messages: [] }), + }); + + await expect(sendRequest(req)).rejects.toThrow( + "Invalid request: messages must be an array and cannot be empty.", ); }); @@ -73,7 +84,7 @@ describe("sendRequest", () => { }); await expect(sendRequest(req)).rejects.toThrow( - "Invalid request: messages must be an array.", + "Invalid request: messages must be an array and cannot be empty.", ); }); @@ -103,13 +114,14 @@ describe("sendRequest", () => { }); it("should throw an error if the user query is too long", async () => { + const messages = [{ role: "user", content: "hello" }]; mockConvertToModelMessages.mockResolvedValue([ { role: "user", content: "a".repeat(10001) }, ]); const req = new Request("http://localhost/api/chat", { method: "POST", - body: JSON.stringify({ messages: [] }), + body: JSON.stringify({ messages }), }); await expect(sendRequest(req)).rejects.toThrow( @@ -118,6 +130,7 @@ describe("sendRequest", () => { }); it("should only send the last 6 messages", async () => { + const messages = [{ role: "user", content: "query" }]; const allMessages = Array.from({ length: 10 }, (_, i) => ({ role: "user", content: `msg ${i}`, @@ -127,7 +140,7 @@ describe("sendRequest", () => { const req = new Request("http://localhost/api/chat", { method: "POST", - body: JSON.stringify({ messages: [] }), + body: JSON.stringify({ messages }), }); await sendRequest(req); diff --git a/src/app/api/chat/request.ts b/src/app/api/chat/request.ts index e0f0f9e..9eaf415 100644 --- a/src/app/api/chat/request.ts +++ b/src/app/api/chat/request.ts @@ -8,8 +8,10 @@ export default async function sendRequest(request: Request) { const body = await request.json(); const { messages } = body; - if (!messages || !Array.isArray(messages)) { - throw new Error("Invalid request: messages must be an array."); + if (!messages || !Array.isArray(messages) || messages.length === 0) { + throw new Error( + "Invalid request: messages must be an array and cannot be empty.", + ); } const modelMessages = await convertToModelMessages(messages);