mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-10 03:29:50 +08:00
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.
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
import { createOpenAI } from "@ai-sdk/openai"
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
|
||||
import {
|
||||
getAIModel,
|
||||
@@ -56,6 +57,9 @@ const ENV_KEYS = [
|
||||
"AI_PROVIDER",
|
||||
"AI_MODEL",
|
||||
"VALIDATION_MODEL",
|
||||
"NEXT_AI_DRAWIO_DESKTOP",
|
||||
"SGLANG_API_KEY",
|
||||
"SGLANG_BASE_URL",
|
||||
]
|
||||
const savedEnv: Record<string, string | undefined> = {}
|
||||
|
||||
@@ -272,3 +276,45 @@ describe("getValidationModel", () => {
|
||||
expect(createOpenRouter).toHaveBeenCalledWith({ apiKey: "env-key" })
|
||||
})
|
||||
})
|
||||
|
||||
describe("whose keys a request uses", () => {
|
||||
it("cleans the base URL like the request does", () => {
|
||||
// "/" and a pasted path clean up to no base URL: the server's Ollama
|
||||
process.env.OLLAMA_API_KEY = "server-ollama-key"
|
||||
expect(usesServerCredentials("ollama", { baseUrl: "/" })).toBe(true)
|
||||
expect(
|
||||
usesServerCredentials("ollama", { baseUrl: "/chat/completions" }),
|
||||
).toBe(true)
|
||||
})
|
||||
|
||||
it("counts the desktop app's keys as the user's own", () => {
|
||||
// Electron passes the user's preset keys as server env vars
|
||||
process.env.NEXT_AI_DRAWIO_DESKTOP = "1"
|
||||
expect(usesServerCredentials("openai", {})).toBe(false)
|
||||
})
|
||||
|
||||
it("sends a user's OpenAI key to the official endpoint", () => {
|
||||
// The SDK would otherwise read the server's OPENAI_BASE_URL
|
||||
process.env.OPENAI_BASE_URL = "https://operator-proxy.example.com/v1"
|
||||
getAIModel({
|
||||
provider: "openai",
|
||||
apiKey: "user-key",
|
||||
modelId: "gpt-5.5",
|
||||
})
|
||||
expect(createOpenAI).toHaveBeenLastCalledWith(
|
||||
expect.objectContaining({
|
||||
apiKey: "user-key",
|
||||
baseURL: "https://api.openai.com/v1",
|
||||
}),
|
||||
)
|
||||
// Still the Responses API, like without a base URL
|
||||
const provider = vi.mocked(createOpenAI).mock.results.at(-1)?.value
|
||||
expect(provider.chat).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("needs a base URL for SGLang instead of using 127.0.0.1", () => {
|
||||
expect(() =>
|
||||
getAIModel({ provider: "sglang", apiKey: "k", modelId: "m" }),
|
||||
).toThrow(/base URL/)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
// @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)
|
||||
})
|
||||
})
|
||||
@@ -110,6 +110,14 @@ describe("classifyLLMError", () => {
|
||||
expect(classifyLLMError(error).code).toBe("model_not_found")
|
||||
})
|
||||
|
||||
it("reads Bedrock's token throttling as a rate limit", () => {
|
||||
const error = apiError(
|
||||
429,
|
||||
"Too many tokens, please wait before trying again.",
|
||||
)
|
||||
expect(classifyLLMError(error).code).toBe("rate_limited")
|
||||
})
|
||||
|
||||
it("names a network error the SDK wrapped", () => {
|
||||
const error = new APICallError({
|
||||
message:
|
||||
@@ -201,6 +209,19 @@ describe("streamErrorText", () => {
|
||||
}
|
||||
})
|
||||
|
||||
it("hides the provider's text on the server's keys", () => {
|
||||
const error = apiError(
|
||||
403,
|
||||
"User: arn:aws:sts::123456789012:assumed-role/app/s is not authorized to perform: bedrock:InvokeModel",
|
||||
)
|
||||
const hidden = JSON.parse(streamErrorText(error, true))
|
||||
expect(hidden.code).toBe("forbidden")
|
||||
expect(hidden.message).not.toMatch(/arn:aws|123456789012/)
|
||||
expect(JSON.parse(streamErrorText(error)).message).toMatch(
|
||||
/not authorized/,
|
||||
)
|
||||
})
|
||||
|
||||
it("classifies a provider error", () => {
|
||||
expect(JSON.parse(streamErrorText(apiError(401, "bad key")))).toEqual({
|
||||
type: "provider",
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
// @vitest-environment node
|
||||
import { mkdtempSync, readFileSync, writeFileSync } from "node:fs"
|
||||
import { tmpdir } from "node:os"
|
||||
import { join } from "node:path"
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest"
|
||||
|
||||
const userData = vi.hoisted(() => ({ dir: "" }))
|
||||
vi.mock("electron", () => ({
|
||||
app: { isPackaged: true, getPath: () => userData.dir },
|
||||
}))
|
||||
|
||||
import { saveServerPort } from "@/electron/main/port-manager"
|
||||
|
||||
const savedPort = () =>
|
||||
JSON.parse(readFileSync(join(userData.dir, "server-port.json"), "utf-8"))
|
||||
.port
|
||||
|
||||
beforeEach(() => {
|
||||
userData.dir = mkdtempSync(join(tmpdir(), "port-manager-"))
|
||||
})
|
||||
|
||||
describe("saveServerPort", () => {
|
||||
it("remembers the port of the first launch", () => {
|
||||
saveServerPort(13370)
|
||||
expect(savedPort()).toBe(13370)
|
||||
})
|
||||
|
||||
it("keeps the remembered port when a launch had to use another", () => {
|
||||
// The app's data lives under the remembered port's origin; going
|
||||
// back to it once it is free brings the chats and settings back
|
||||
writeFileSync(
|
||||
join(userData.dir, "server-port.json"),
|
||||
JSON.stringify({ port: 61337 }),
|
||||
)
|
||||
saveServerPort(13371)
|
||||
expect(savedPort()).toBe(61337)
|
||||
})
|
||||
})
|
||||
@@ -103,6 +103,29 @@ describe("listProviderModels", () => {
|
||||
])
|
||||
})
|
||||
|
||||
it("lists Ollama from where chat goes without a base URL", async () => {
|
||||
const { fn, calls } = answer({ models: [{ name: "llama3.2" }] })
|
||||
process.env.OLLAMA_BASE_URL = "http://ollama.internal:11434"
|
||||
try {
|
||||
await listProviderModels("ollama", {}, fn)
|
||||
} finally {
|
||||
delete process.env.OLLAMA_BASE_URL
|
||||
}
|
||||
await listProviderModels("ollama", {}, fn)
|
||||
expect(calls.map((c) => c.url)).toEqual([
|
||||
"http://ollama.internal:11434/api/tags",
|
||||
"http://127.0.0.1:11434/api/tags",
|
||||
])
|
||||
})
|
||||
|
||||
it("does not use SGLang's local address as a default", async () => {
|
||||
const { fn, calls } = answer({ data: [] })
|
||||
await expect(
|
||||
listProviderModels("sglang", { apiKey: "k" }, fn),
|
||||
).rejects.toThrow(/base URL/)
|
||||
expect(calls).toHaveLength(0)
|
||||
})
|
||||
|
||||
it("turns a failed request into an error with its status", async () => {
|
||||
const { fn } = answer({ error: "bad key" }, 401)
|
||||
await expect(
|
||||
|
||||
@@ -74,6 +74,50 @@ describe("POST /api/validate-model", () => {
|
||||
expect(typeof data.responseTime).toBe("number")
|
||||
})
|
||||
|
||||
it("reports a model that did not answer in time", async () => {
|
||||
// The 15 s timeout has fired: the SDK ends the stream with an
|
||||
// abort part instead of throwing
|
||||
const timedOut = AbortSignal.abort(
|
||||
new DOMException("The operation timed out.", "TimeoutError"),
|
||||
)
|
||||
const timeout = vi
|
||||
.spyOn(AbortSignal, "timeout")
|
||||
.mockReturnValue(timedOut)
|
||||
vi.stubGlobal(
|
||||
"fetch",
|
||||
vi.fn(async () => {
|
||||
throw timedOut.reason
|
||||
}),
|
||||
)
|
||||
try {
|
||||
const data = await testGlm()
|
||||
expect(data.valid).toBe(false)
|
||||
expect(data.code).toBe("timeout")
|
||||
} finally {
|
||||
timeout.mockRestore()
|
||||
}
|
||||
})
|
||||
|
||||
it("does not run on the server's keys", async () => {
|
||||
process.env.OLLAMA_API_KEY = "server-ollama-key"
|
||||
try {
|
||||
const res = await validateModel(
|
||||
new Request("http://localhost/api/validate-model", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
provider: "ollama",
|
||||
modelId: "any-cloud-model",
|
||||
}),
|
||||
}),
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
expect((await res.json()).error).toMatch(/API key/)
|
||||
} finally {
|
||||
delete process.env.OLLAMA_API_KEY
|
||||
}
|
||||
})
|
||||
|
||||
it("warns when the model answers without a tool call", async () => {
|
||||
streamReply({ role: "assistant", content: "OK" })
|
||||
const data = await testGlm()
|
||||
|
||||
Reference in New Issue
Block a user