i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-12 20:49:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(chat): use the AI SDK's own mechanisms for quota, retries, stops and notices

Browse Source 
Server
- Record quota tokens in onStepFinish and await the write, so a step that
  fails or is stopped later still counts; drop the onFinish/onAbort copies
- validate-diagram and the connection test end their model call with the
  request (req.signal) and use the SDK's timeout option; validate-diagram
  reports to Langfuse (functionId "validate-diagram")
- The connection test no longer counts an invalid tool call as support
- Validate incoming messages with safeValidateUIMessages (400 otherwise),
  drop client system messages explicitly, prune empty model messages with
  pruneMessages
- Chat tools move to lib/chat-tools.ts with an output schema for the
  client-side tools

Client
- Every automatic retry sends the canvas as it is then
  (prepareSendMessagesRequest): a drawing is committed before its
  screenshot check can fail, so the model must go on from it
- A notice when the output budget ran out before any tool call
  (finishReason "length")
- Error notices are assistant messages with a data-notice part instead of
  role "system" messages (lib/chat-message.ts: ChatUIMessage with metadata,
  data part and InferUITools types)
- regenerate, edit and retry use sendMessage({ messageId }); isResending
  and flushSync go
- isToolUIPart/getToolName instead of slicing "tool-"; ToolPartLike derives
  from the SDK's part types; addToolOutput typed with ChatAddToolOutputFunction

Tests: step failure after a finished step, abort of the screenshot check,
invalid tool call in the connection test, malformed messages.
This commit is contained in:
dayuan.jiang
2026-10-12 12:36:45 +09:00
parent 60105ddb51
commit 92e785bcb5
16 changed files with 578 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/chat-route-abort.test.ts
+101 -79
View File
@@ -28,60 +28,89 @@ const sse = (chunks: object[], end = true) =>
chunks.map((c) => `data: ${JSON.stringify(c)}\n\n`).join("") +
(end ? "data: [DONE]\n\n" : "")
// Step 1: the model asks for a shape library (a tool run on the server) and
// reports its usage, so the chat makes a second request
const firstStep = () =>
new Response(
sse([
{
id: "c1",
choices: [
{
index: 0,
delta: {
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: {
name: "get_shape_library",
arguments: '{"library":"aws4"}',
},
},
],
},
finish_reason: null,
},
],
},
{
id: "c1",
choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }],
usage: { prompt_tokens: 1200, completion_tokens: 30 },
},
]),
{ headers: { "content-type": "text/event-stream" } },
)
const send = (signal?: AbortSignal) =>
chat(
new Request("http://localhost/api/chat", {
method: "POST",
signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
// The server's own network: counted
"x-ai-provider": "glm",
"x-ai-base-url": "http://127.0.0.1:9000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "glm-5",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw AWS" }],
},
],
xml: "",
}),
}),
)
const drain = async (res: Response) => {
const reader = res.body?.getReader()
while (
reader &&
!(await reader.read().catch(() => ({ done: true }))).done
) {
// drain
}
}
describe("a request stopped after a finished step", () => {
it("counts that step's tokens", async () => {
// Step 1 asks for a shape library (run on the server) and reports
// its usage; step 2 never ends, and the user stops
// Step 2 never ends, and the user stops
let call = 0
vi.stubGlobal(
"fetch",
vi.fn(async (_url: string, init?: RequestInit) => {
call++
if (call === 1) {
return new Response(
sse([
{
id: "c1",
choices: [
{
index: 0,
delta: {
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: {
name: "get_shape_library",
arguments:
'{"library":"aws4"}',
},
},
],
},
finish_reason: null,
},
],
},
{
id: "c1",
choices: [
{
index: 0,
delta: {},
finish_reason: "tool_calls",
},
],
usage: {
prompt_tokens: 1200,
completion_tokens: 30,
},
},
]),
{ headers: { "content-type": "text/event-stream" } },
)
}
if (call === 1) return firstStep()
// Never ends, until the request is aborted (as fetch does)
const body = new ReadableStream({
start(controller) {
@@ -98,42 +127,35 @@ describe("a request stopped after a finished step", () => {
}),
)
const stop = new AbortController()
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
signal: stop.signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
// The server's own network: counted
"x-ai-provider": "glm",
"x-ai-base-url": "http://127.0.0.1:9000/v1",
"x-ai-api-key": "dummy",
"x-ai-model": "glm-5",
},
body: JSON.stringify({
messages: [
{
id: "u1",
role: "user",
parts: [{ type: "text", text: "Draw AWS" }],
},
],
xml: "",
}),
}),
)
const reader = res.body?.getReader()
const res = await send(stop.signal)
const draining = drain(res)
// Read until the second step has started
await vi.waitFor(() => expect(call).toBe(2), { timeout: 3000 })
stop.abort()
// The answer stream ends; the SDK handles the stop as it is read
while (
reader &&
!(await reader.read().catch(() => ({ done: true }))).done
) {
// drain
}
await draining
await vi.waitFor(() => expect(quota.recorded).toEqual([1230]))
})
})
describe("a request that fails after a finished step", () => {
it("counts that step's tokens", async () => {
// Step 2 is refused by the provider (a status the SDK does not
// retry): the stream ends with an error
let call = 0
vi.stubGlobal(
"fetch",
vi.fn(async () => {
call++
if (call === 1) return firstStep()
return new Response('{"error":{"message":"bad request"}}', {
status: 400,
headers: { "content-type": "application/json" },
})
}),
)
await drain(await send())
expect(call).toBe(2)
expect(quota.recorded).toEqual([1230])
})
})
tests/unit/chat-route-quota.test.ts
+15
View File
@@ -205,6 +205,21 @@ describe("request checks", () => {
expect(await res.text()).toMatch(/Invalid AWS region/)
expect(fetch).not.toHaveBeenCalled()
})
it("refuses messages that are not UI messages", async () => {
const res = await chat(
new Request("http://localhost/api/chat", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
messages: [{ id: "u1", role: "user", text: "no parts" }],
xml: "",
}),
}),
)
expect(res.status).toBe(400)
expect(fetch).not.toHaveBeenCalled()
})
})
describe("server model allowlist", () => {
tests/unit/validate-diagram-route.test.ts
+40 -11
View File
@@ -15,19 +15,32 @@ const RESULT = {
suggestions: ["Move box B to the right"],
}
// A vision model that answers with the result as JSON text
// A vision model that answers with the result as JSON text; told to wait,
// it answers only once its call is aborted
const model = vi.hoisted(() => ({ waitForAbort: false }))
vi.mock("@/lib/ai-providers", () => ({
getValidationModel: () =>
new MockLanguageModelV3({
doGenerate: (async () => ({
content: [{ type: "text", text: JSON.stringify(RESULT) }],
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
warnings: [],
})) as any,
doGenerate: (async (options: { abortSignal?: AbortSignal }) => {
if (model.waitForAbort) {
const signal = options.abortSignal
await new Promise((_, reject) => {
if (signal?.aborted) reject(signal.reason)
signal?.addEventListener("abort", () =>
reject(signal.reason),
)
})
}
return {
content: [{ type: "text", text: JSON.stringify(RESULT) }],
finishReason: { unified: "stop", raw: "stop" },
usage: {
inputTokens: { total: 1 },
outputTokens: { total: 1 },
},
warnings: [],
}
}) as any,
}),
}))
@@ -61,9 +74,10 @@ vi.mock("@/lib/dynamo-quota-manager", () => ({
},
}))
const post = () =>
const post = (signal?: AbortSignal) =>
new Request("http://localhost/api/validate-diagram", {
method: "POST",
signal,
headers: {
"Content-Type": "application/json",
"x-forwarded-for": "203.0.113.7",
@@ -72,6 +86,7 @@ const post = () =>
})
afterEach(() => {
model.waitForAbort = false
quota.enabled = false
quota.allowed = true
quota.checks = []
@@ -106,6 +121,20 @@ describe("POST /api/validate-diagram", () => {
expect(await res.json()).toEqual(RESULT)
})
it("ends its model call when the request is aborted", async () => {
// The user pressed Stop: the browser aborts the request, and the
// model call must not run on (and be paid for)
model.waitForAbort = true
const stop = new AbortController()
const pending = validateDiagram(post(stop.signal))
stop.abort()
expect(await (await pending).json()).toEqual({
valid: true,
issues: [],
suggestions: [],
})
})
it("answers valid when the check is turned off", async () => {
vi.stubEnv("ENABLE_VLM_VALIDATION", "false")
const res = await validateDiagram(post())
tests/unit/validate-model-route.test.ts
+18
View File
@@ -141,6 +141,24 @@ describe("POST /api/validate-model", () => {
expect(data.valid).toBe(true)
expect(data.warning).toMatch(/without calling a tool/)
})
it("warns when the model calls a tool that does not exist", async () => {
// The SDK passes the call on, marked invalid
streamReply({
role: "assistant",
tool_calls: [
{
index: 0,
id: "call_1",
type: "function",
function: { name: "pong", arguments: "{}" },
},
],
})
const data = await testGlm()
expect(data.valid).toBe(true)
expect(data.warning).toMatch(/without calling a tool/)
})
})
describe("chat requests to a client base URL", () => {