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

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
+11 -5
View File
@@ -11,6 +11,7 @@ import {
isQuotaEnabled,
recordTokenUsage,
} from "@/lib/dynamo-quota-manager"
import { getTelemetryConfig } from "@/lib/langfuse"
import { getUserIdFromRequest } from "@/lib/user-id"
import { VALIDATION_SYSTEM_PROMPT } from "@/lib/validation-prompts"
import {
@@ -146,13 +147,18 @@ export async function POST(req: Request): Promise<Response> {
},
],
maxOutputTokens: 1024,
abortSignal: AbortSignal.timeout(timeout),
// Ends with the request: the user pressed Stop, or left
abortSignal: req.signal,
timeout,
// Its usage and cost show in Langfuse next to the chat's
experimental_telemetry: getTelemetryConfig({
functionId: "validate-diagram",
sessionId,
userId,
}),
})
if (countsQuota) {
await recordTokenUsage(
userId,
(totalUsage.inputTokens || 0) + (totalUsage.outputTokens || 0),
)
await recordTokenUsage(userId, totalUsage.totalTokens ?? 0)
}
if (sessionId) {
console.log(