Files
next-ai-draw-io/components/chat/tool-activity.tsx
T
Dayuan Jiang 99ced88e5b 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
2026-10-12 12:58:52 +09:00

366 lines
13 KiB
TypeScript

import { getToolName } from "ai"
import {
AlertCircle,
Check,
ChevronDown,
Copy,
Library,
PenTool,
} from "lucide-react"
import { useState } from "react"
import { Shimmer } from "@/components/ai-elements/shimmer"
import { VersionCard as VersionCardView } from "@/components/canvas/version-card"
import { CodeBlock } from "@/components/code-block"
import { useDictionary } from "@/hooks/use-dictionary"
import { useCopy } from "@/lib/clipboard"
import { cn, isMxCellXmlComplete } from "@/lib/utils"
import { type DiagramVersion, useVersionsStore } from "@/stores/versions-store"
import { useChatEngine } from "./chat-engine"
import type { DiagramOperation, ToolPartLike } from "./types"
export function ToolInputDetails({ part }: { part: ToolPartLike }) {
const { input, state } = part
if (!input || typeof input !== "object") return null
if (typeof input.xml === "string" && input.xml) {
// Plain text while streaming; highlighting re-runs on every chunk
return state === "input-streaming" || state === "input-available" ? (
<pre className="max-h-48 overflow-auto whitespace-pre-wrap break-all font-mono text-[11px] leading-relaxed text-muted-foreground scrollbar-thin">
{input.xml}
</pre>
) : (
<CodeBlock code={input.xml} language="xml" />
)
}
if (Array.isArray(input.operations)) {
// Streamed or invalid input can hold anything: show only what React
// can render (an object in place of a string would crash the chat)
const shown = (input.operations as DiagramOperation[]).filter(
(op) =>
typeof (op as { operation?: unknown })?.operation === "string",
)
const text = (value: unknown) =>
typeof value === "string" ? value : ""
return (
<div className="space-y-2">
{shown.map((op, index) => (
<div
key={`${op.operation}-${text(op.cell_id)}-${index}`}
className="overflow-hidden rounded-lg border border-border bg-card"
>
<div className="flex items-center gap-2 border-b border-border px-2.5 py-1 text-[11px]">
<span className="font-medium text-foreground">
{op.operation}
</span>
<span className="text-muted-foreground">
{text(op.cell_id)}
</span>
</div>
{text(op.new_xml) && (
<pre className="overflow-x-auto whitespace-pre-wrap break-all px-2.5 py-1.5 font-mono text-[11px] text-muted-foreground">
{text(op.new_xml)}
</pre>
)}
</div>
))}
</div>
)
}
if (Object.keys(input).length === 0) return null
return <CodeBlock code={JSON.stringify(input, null, 2)} language="json" />
}
function toolText(part: ToolPartLike): string {
const { input } = part
if (!input || typeof input !== "object") return ""
if (typeof input.xml === "string" && input.xml) return input.xml
if (Array.isArray(input.operations)) {
return JSON.stringify(input.operations, null, 2)
}
return JSON.stringify(input, null, 2)
}
/**
* Card for one diagram version produced by a tool call: the shared card
* with a link to what the AI wrote (shown in a panel under the card)
*/
function VersionCard({
part,
version,
isLatest,
}: {
part: ToolPartLike
version: DiagramVersion
isLatest: boolean
}) {
const dict = useDictionary()
const t = dict.versions
const [showCode, setShowCode] = useState(false)
const { copied, copy } = useCopy()
const isEdit = part.type === "tool-edit_diagram"
return (
<VersionCardView
version={version}
isLatest={isLatest}
label={isEdit ? t.edited : t.created}
data-tool-state={part.state}
// A text link, so the card's only button is the action (undo
// or restore); it wraps as one unit with its separator
afterSummary={
<span className="whitespace-nowrap">
<span aria-hidden> · </span>
<button
type="button"
onClick={() => setShowCode((v) => !v)}
className="underline decoration-foreground/30 underline-offset-2 hover:text-foreground hover:decoration-foreground"
>
{showCode ? t.hideCode : t.showCode}
</button>
</span>
}
>
{showCode && (
<div className="relative border-t border-border bg-surface-1 px-3 py-2.5">
<button
type="button"
onClick={() => copy(toolText(part))}
className="absolute top-2 right-2 z-10 inline-flex size-7 items-center justify-center rounded-lg bg-card text-muted-foreground shadow-float hover:text-foreground"
aria-label={dict.chat.copyResponse}
>
{copied ? (
<Check className="size-3.5" />
) : (
<Copy className="size-3.5" />
)}
</button>
<ToolInputDetails part={part} />
</div>
)}
</VersionCardView>
)
}
/** Compact row for tool calls that are running, failed or produced no version */
function ToolRow({
part,
icon,
label,
tone = "default",
detail,
}: {
part: ToolPartLike
icon: React.ReactNode
label: React.ReactNode
tone?: "default" | "error" | "warning"
detail?: string
}) {
const dict = useDictionary()
const [open, setOpen] = useState(false)
const { copied, copy } = useCopy()
const hasInput = !!part.input && Object.keys(part.input).length > 0
return (
<div
className={cn(
"overflow-hidden rounded-xl border border-border bg-card",
tone === "error" && "border-destructive/25",
)}
data-testid="tool-row"
data-tool-name={getToolName(part)}
data-tool-state={part.state}
>
<button
type="button"
onClick={() => setOpen((v) => !v)}
disabled={!hasInput && !detail}
className="flex w-full items-center gap-2.5 px-3 py-2.5 text-left"
>
<span
className={cn(
"inline-flex size-7 shrink-0 items-center justify-center rounded-lg bg-muted text-muted-foreground [&_svg]:size-4",
tone === "error" &&
"bg-destructive/10 text-destructive",
tone === "warning" && "bg-marker-soft text-marker-ink",
)}
>
{icon}
</span>
<span className="min-w-0 flex-1 truncate text-[13px] text-foreground">
{label}
</span>
{(hasInput || detail) && (
<ChevronDown
className={cn(
"size-4 shrink-0 text-muted-foreground transition-transform",
open && "rotate-180",
)}
/>
)}
</button>
{open && (
<div
className={cn(
"relative space-y-2 border-t border-border bg-surface-1 px-3 py-2.5",
hasInput && "pr-11",
)}
>
{hasInput && (
<button
type="button"
onClick={() => copy(toolText(part))}
className="absolute top-2 right-2 z-10 inline-flex size-7 items-center justify-center rounded-lg bg-card text-muted-foreground shadow-float hover:text-foreground"
aria-label={dict.chat.copyResponse}
>
{copied ? (
<Check className="size-3.5" />
) : (
<Copy className="size-3.5" />
)}
</button>
)}
{detail && (
<p
className={cn(
"whitespace-pre-wrap break-words text-xs",
tone === "error"
? "text-destructive"
: "text-muted-foreground",
)}
>
{detail}
</p>
)}
{hasInput && <ToolInputDetails part={part} />}
</div>
)}
</div>
)
}
export function ToolActivity({
part,
isLastMessage,
}: {
part: ToolPartLike
isLastMessage: boolean
}) {
const dict = useDictionary()
const t = dict.versions
const engine = useChatEngine()
const versions = useVersionsStore((s) => s.versions)
const toolName = getToolName(part)
const { state, errorText, output } = part
const versionIndex = versions.findIndex(
(v) => v.toolCallId === part.toolCallId,
)
if (state === "output-available" && versionIndex >= 0) {
return (
<VersionCard
part={part}
version={versions[versionIndex]}
isLatest={versionIndex === versions.length - 1}
/>
)
}
// The answer ended before the call did (stopped, network error)
const running = isLastMessage && engine.isBusy
if (toolName === "get_shape_library" && state !== "output-error") {
const isDone = state === "output-available"
return (
<ToolRow
part={part}
icon={<Library />}
label={
isDone ? (
t.libraryLoaded
) : running ? (
<Shimmer as="span">{t.loadingLibrary}</Shimmer>
) : (
t.stopped
)
}
detail={
isDone && typeof output === "string"
? `${output.slice(0, 600)}${output.length > 600 ? "…" : ""}`
: undefined
}
/>
)
}
// Stop turns a call still streaming into "Stopped by user" (right
// away, or when the answer has ended for one not on screen yet): one
// left like this was cut off (a closed connection, an error, a closed
// tab)
if (state === "input-streaming" || state === "input-available") {
const failed = !running && !(isLastMessage && engine.stoppedRef.current)
return (
<ToolRow
part={part}
icon={running ? <PenTool /> : <AlertCircle />}
tone={failed ? "error" : "default"}
label={
running ? (
<Shimmer as="span">
{toolName === "edit_diagram"
? t.editing
: t.drawing}
</Shimmer>
) : failed ? (
t.failed
) : (
t.stopped
)
}
/>
)
}
if (state === "output-error") {
const stopped = errorText === "Stopped by user"
// Incomplete XML means the output hit the length limit. Without an
// input the JSON was broken (the server repairs JSON cut short by
// the limit).
const truncated =
!stopped &&
(toolName === "display_diagram" || toolName === "append_diagram") &&
typeof part.input?.xml === "string" &&
!isMxCellXmlComplete(part.input.xml)
return (
<ToolRow
part={part}
icon={<AlertCircle />}
tone={stopped ? "default" : truncated ? "warning" : "error"}
label={
stopped
? t.stopped
: truncated
? dict.tools.truncated
: isLastMessage && engine.isBusy
? t.failedRetrying
: t.failed
}
detail={
stopped
? undefined
: truncated
? t.truncatedHint
: (errorText ??
(typeof output === "string" ? output : ""))
}
/>
)
}
// Finished without a version (sessions saved before versions existed)
return (
<ToolRow
part={part}
icon={<PenTool />}
label={toolName === "edit_diagram" ? t.edited : t.created}
/>
)
}