mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-12 20:49:53 +08:00
fix(chat): use the AI SDK's own mechanisms for quota, retries, stops and notices (#987)
## Problem
A review of how the app uses the Vercel AI SDK (two Claude reviewers and Codex, each finding checked against the SDK's type definitions and source) found five behaviours that are wrong and a set of places where the app re-implements what the SDK already provides:
1. Quota tokens were recorded in `onFinish`, which the SDK does not call when a later step fails; only a user stop was handled (`onAbort`), and the write was not awaited.
2. An automatic retry after a tool error re-sent the request body as it was when the message was sent, so the model was told the "authoritative" canvas XML was the diagram from before the turn, although `display_diagram` commits a drawing before its screenshot check can fail.
3. The screenshot check and the connection test ran their model call on `AbortSignal.timeout` only, so pressing Stop (or closing the dialog) did not stop a paid model call.
4. The client ignored `finishReason`; a reasoning model that spent the whole output budget before calling a tool left the user with no diagram and no message.
5. The connection test counted an `invalid` tool call (a tool that does not exist, unreadable input) as proof of tool support.
## Change
- `onStepFinish` records each step's `usage.totalTokens` and awaits the write; the copies in `onFinish` and `onAbort` are gone.
- `DefaultChatTransport.prepareSendMessagesRequest` sends the current canvas page XML on every request after the first of a turn; the first keeps the exported (or snapshot) XML.
- `validate-diagram` and the connection test pass `req.signal` and use the SDK's `timeout` option; `validate-diagram` reports to Langfuse under `functionId: "validate-diagram"`.
- The connection test only counts a valid tool call.
- The client's `onFinish` shows a notice (with Retry) when `finishReason === "length"` and no tool was called.
- The server validates incoming messages with `safeValidateUIMessages` (400 otherwise), drops client system messages explicitly and prunes empty model messages with `pruneMessages`.
- The chat's tools live in `lib/chat-tools.ts` with an `outputSchema` for the client-side ones; `lib/chat-message.ts` defines `ChatUIMessage` (metadata, a `data-notice` part, `InferUITools`) so `useChat`, `addToolOutput` and the tool handlers are typed.
- Error notices are assistant messages with a `data-notice` part instead of role `system` UI messages.
- Regenerate, edit and retry use `sendMessage({ messageId })`; `isResending` and `flushSync` are gone.
- `isToolUIPart` / `getToolName` replace slicing `"tool-"` off part types; `ToolPartLike` derives from the SDK's part types.
Net: 578 insertions, 439 deletions over 16 files.
## Testing
- New unit tests: a step that fails after a finished step still counts that step's tokens; the screenshot check's model call ends when the request is aborted; an invalid tool call in the connection test gives the warning; malformed messages get a 400. The first and third confirm SDK behaviour the review relied on.
- `tsc`, `biome ci` (0 errors), root `vitest run` 723 passed on the branch rebased onto #986, `next build`: clean
- Playwright, 10 chat-related specs on the dev server: 199 passed, 0 failed, both before and after the rebase onto #986
This commit is contained in:
Dayuan Jiang
@@ -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])
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
Reference in New Issue
Block a user