Files
next-ai-draw-io/tests/unit/chat-route-quota.test.ts
T
dayuan.jiang 0855b35ff2 fix(server): count quota by the key actually used, and more review fixes
Found by the PR review, each with a test that failed first:
- Quota: any key header skipped it, even one the provider never reads
  (x-aws-access-key-id with OpenAI), so a request ran on the server's
  key without being counted. The check now runs after the model is
  resolved and uses usesServerCredentials. On main already.
- usesServerCredentials read the raw base URL; "/" cleans up to none, so
  an Ollama request ran on the server's key past the server-model check.
- SGLang's default 127.0.0.1:8000 only fills the settings form. Chat and
  the model list used it as a real address, so the server called its own
  machine even with private URLs blocked. Now a base URL is required.
- With a user's OpenAI key and no base URL, the SDK read the server's
  OPENAI_BASE_URL. The official endpoint is now passed. On main already.
- The Test button refused nothing on the server's keys (Ollama Cloud),
  and a 15 s timeout reported "connected, no tool call".
- The model list for Ollama without a base URL came from ollama.com while
  chat went to the server's Ollama.
- Bedrock's "Too many tokens, please wait" counted as context too long.
- On the server's keys the provider's error text stays in the server log;
  it can name the server's AWS account, role or internal hosts.
- Desktop app: the preset keys are the user's own (NEXT_AI_DRAWIO_DESKTOP),
  so Max Output Tokens can be raised and keyless models in settings work
  again. A launch that found the remembered port taken no longer replaces
  it, which hid the user's chats and settings for good.
2026-10-04 23:04:21 +09:00

92 lines
2.6 KiB
TypeScript

// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
// Quota on, and every check answers that the daily limit is used up
const quota = vi.hoisted(() => ({ checks: 0 }))
vi.mock("@/lib/dynamo-quota-manager", () => ({
isQuotaEnabled: () => true,
checkAndIncrementRequest: async () => {
quota.checks++
return {
allowed: false,
error: "Daily limit reached",
type: "request",
used: 10,
limit: 10,
}
},
recordTokenUsage: async () => {},
}))
import { POST as chat } from "@/app/api/chat/route"
const ENV = ["AI_PROVIDER", "AI_MODEL", "OPENAI_API_KEY"]
const saved: Record<string, string | undefined> = {}
beforeEach(() => {
for (const k of ENV) saved[k] = process.env[k]
process.env.AI_PROVIDER = "openai"
process.env.AI_MODEL = "gpt-5.5"
process.env.OPENAI_API_KEY = "server-key"
quota.checks = 0
// No request may reach a provider
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw new Error("no network in tests")
}),
)
})
afterEach(() => {
for (const k of ENV) {
if (saved[k] === undefined) delete process.env[k]
else process.env[k] = saved[k]
}
vi.unstubAllGlobals()
})
const send = (headers: Record<string, string>) =>
chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
...headers,
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw two boxes" }],
},
],
xml: "",
}),
}),
)
describe("chat quota", () => {
it("counts a request whose key header the provider never reads", async () => {
// OpenAI ignores the AWS key, so this runs on the server's key
const res = await send({
"x-ai-provider": "openai",
"x-aws-access-key-id": "x",
})
expect(res.status).toBe(429)
expect(quota.checks).toBe(1)
})
it("does not count a request on the user's own key", async () => {
const res = await send({
"x-ai-provider": "openai",
"x-ai-api-key": "user-key",
"x-ai-model": "gpt-5.5",
})
expect(res.status).not.toBe(429)
expect(quota.checks).toBe(0)
})
})