handle data parts in AIChat

.changeset/cuddly-chicken-argue.md

@@ -0,0 +1,13 @@
+---
+"@cloudflare/ai-chat": patch
+---
+
+Add support for `data-*` stream parts (developer-defined typed JSON blobs) in the shared message builder and client hook.
+
+**Data part handling:**
+
+`applyChunkToParts` now handles `data-*` prefixed chunk types, covering both server persistence and client reconstruction (stream resume, cross-tab broadcast). Transient parts (`transient: true`) are broadcast to connected clients but excluded from `message.parts` and SQLite persistence. Non-transient parts support reconciliation by type+id — a second chunk with the same type and id updates the existing part's data in-place instead of appending a duplicate.
+
+**`onData` callback forwarding:**
+
+`useAgentChat` now invokes the `onData` callback for `data-*` chunks on the stream resumption and cross-tab broadcast codepaths, which bypass the AI SDK's internal pipeline. For new messages sent via the transport, the AI SDK already invokes `onData` internally. This is the correct way to consume transient data parts on the client since they are not added to `message.parts`.

docs/AGENTS.md

@@ -23,13 +23,7 @@ The fourth Diátaxis type — **explanation** ("understand why") — lives in `/
 
 ## Upstream sync
 
-Changes to `docs/` trigger a CI workflow (`.github/workflows/sync-docs.yml`) that:
-
-1. Diffs the PR against main to find changed doc files
-2. Runs Claude Code to adapt the content for `cloudflare/cloudflare-docs` conventions
-3. Creates or updates a PR in the `cloudflare/cloudflare-docs` repo
-
-This means your markdown here gets transformed at sync time — the sync handles Cloudflare-specific components (`WranglerConfig`, `PackageManagers`, `TypeScriptExample`, etc.), link format conversion, and style adjustments. Write clean, accurate markdown here; the sync workflow handles the rest.
+There is no automated sync workflow. Changes here must be manually ported to `cloudflare/cloudflare-docs` (the `src/content/docs/agents/` folder). When updating docs here, also update the corresponding `.mdx` file in cloudflare-docs — or at minimum note the discrepancy in `docs-to-upstream.md`.
 
 ### Tracking file
 

docs/chat-agents.md

@@ -17,6 +17,7 @@ Built on the [AI SDK](https://ai-sdk.dev) and Cloudflare Durable Objects, you ge
 - **Resumable streaming** — disconnected clients resume mid-stream without data loss
 - **Real-time sync** — messages broadcast to all connected clients via WebSocket
 - **Tool support** — server-side, client-side, and human-in-the-loop tool patterns
+- **Data parts** — attach typed JSON (citations, progress, usage) to messages alongside text
 - **Row size protection** — automatic compaction when messages approach SQLite limits
 
 ## Quick Start
@@ -515,6 +516,128 @@ const { messages, sendMessage } = useAgentChat({
 });
 ```
 
+## Data Parts
+
+Data parts let you attach typed JSON to messages alongside text — progress indicators, source citations, token usage, or any structured data your UI needs.
+
+### Writing Data Parts (Server)
+
+Use `createUIMessageStream` with `writer.write()` to send data parts from the server:
+
+```ts
+import {
+  streamText,
+  convertToModelMessages,
+  createUIMessageStream,
+  createUIMessageStreamResponse
+} from "ai";
+
+export class ChatAgent extends AIChatAgent {
+  async onChatMessage() {
+    const workersai = createWorkersAI({ binding: this.env.AI });
+
+    const stream = createUIMessageStream({
+      execute: async ({ writer }) => {
+        const result = streamText({
+          model: workersai("@cf/zai-org/glm-4.7-flash"),
+          messages: await convertToModelMessages(this.messages)
+        });
+
+        // Merge the LLM stream
+        writer.merge(result.toUIMessageStream());
+
+        // Write a data part — persisted to message.parts
+        writer.write({
+          type: "data-sources",
+          id: "src-1",
+          data: { query: "agents", status: "searching", results: [] }
+        });
+
+        // Later: update the same part in-place (same type + id)
+        writer.write({
+          type: "data-sources",
+          id: "src-1",
+          data: {
+            query: "agents",
+            status: "found",
+            results: ["Agents SDK docs", "Durable Objects guide"]
+          }
+        });
+      }
+    });
+
+    return createUIMessageStreamResponse({ stream });
+  }
+}
+```
+
+### Three Patterns
+
+| Pattern            | How                                              | Persisted? | Use case                              |
+| ------------------ | ------------------------------------------------ | ---------- | ------------------------------------- |
+| **Reconciliation** | Same `type` + `id` → updates in-place            | Yes        | Progressive state (searching → found) |
+| **Append**         | No `id`, or different `id` → appends             | Yes        | Log entries, multiple citations       |
+| **Transient**      | `transient: true` → not added to `message.parts` | No         | Ephemeral status (thinking indicator) |
+
+Transient parts are broadcast to connected clients in real time but excluded from SQLite persistence and `message.parts`. Use the `onData` callback to consume them.
+
+### Reading Data Parts (Client)
+
+Non-transient data parts appear in `message.parts`. Use the `UIMessage` generic to type them:
+
+```ts
+import { useAgentChat } from "@cloudflare/ai-chat/react";
+import type { UIMessage } from "ai";
+
+type ChatMessage = UIMessage<
+  unknown,
+  {
+    sources: { query: string; status: string; results: string[] };
+    usage: { model: string; inputTokens: number; outputTokens: number };
+  }
+>;
+
+const { messages } = useAgentChat<unknown, ChatMessage>({ agent });
+
+// Typed access — no casts needed
+for (const msg of messages) {
+  for (const part of msg.parts) {
+    if (part.type === "data-sources") {
+      console.log(part.data.results); // string[]
+    }
+  }
+}
+```
+
+### Transient Parts with `onData`
+
+Transient data parts are not in `message.parts`. Use the `onData` callback instead:
+
+```ts
+const [thinking, setThinking] = useState(false);
+
+const { messages } = useAgentChat<unknown, ChatMessage>({
+  agent,
+  onData(part) {
+    if (part.type === "data-thinking") {
+      setThinking(true);
+    }
+  }
+});
+```
+
+On the server, write transient parts with `transient: true`:
+
+```ts
+writer.write({
+  transient: true,
+  type: "data-thinking",
+  data: { model: "glm-4.7-flash", startedAt: new Date().toISOString() }
+});
+```
+
+`onData` fires on all code paths — new messages, stream resumption, and cross-tab broadcasts.
+
 ## Resumable Streaming
 
 Streams automatically resume when a client disconnects and reconnects. No configuration is needed — it works out of the box.

examples/resumable-stream-chat/src/client.tsx

@@ -11,15 +11,116 @@ import {
 import {
   PaperPlaneRightIcon,
   TrashIcon,
-  ArrowClockwiseIcon
+  ArrowClockwiseIcon,
+  MagnifyingGlassIcon,
+  BrainIcon,
+  ChartBarIcon
 } from "@phosphor-icons/react";
 import type { UIMessage } from "ai";
 
-/** Extract plain text from a UIMessage's parts. */
-function getMessageText(message: UIMessage): string {
+// ── Typed data parts ──────────────────────────────────────────
+
+type SourcesData = {
+  query: string;
+  status: "searching" | "found";
+  results: string[];
+};
+
+type ThinkingData = {
+  model: string;
+  startedAt: string;
+};
+
+type UsageData = {
+  model: string;
+  inputTokens: number;
+  outputTokens: number;
+  latencyMs: number;
+};
+
+/** Custom message type with typed data parts. */
+type ChatMessage = UIMessage<
+  unknown,
+  {
+    sources: SourcesData;
+    thinking: ThinkingData;
+    usage: UsageData;
+  }
+>;
+
+// ── Data part renderers ─────────────────────────────────────────────
+
+function SourcesPart({
+  data,
+  isStreaming
+}: {
+  data: SourcesData;
+  isStreaming: boolean;
+}) {
+  if (data.status === "searching") {
+    return (
+      <div className="flex items-center gap-2 text-xs text-kumo-subtle py-1.5">
+        <MagnifyingGlassIcon
+          size={14}
+          className={isStreaming ? "animate-pulse-dot" : ""}
+        />
+        <span>Searching for &ldquo;{data.query}&rdquo;&hellip;</span>
+      </div>
+    );
+  }
+
+  return (
+    <div className="text-xs border border-kumo-line rounded-lg p-2.5 mb-2">
+      <div className="flex items-center gap-1.5 text-kumo-subtle mb-1.5">
+        <MagnifyingGlassIcon size={12} />
+        <span className="font-medium">Sources</span>
+      </div>
+      <ul className="space-y-0.5">
+        {data.results.map((source) => (
+          <li
+            key={source}
+            className="text-kumo-default pl-3 relative before:content-['·'] before:absolute before:left-0 before:text-kumo-subtle"
+          >
+            {source}
+          </li>
+        ))}
+      </ul>
+    </div>
+  );
+}
+
+function ThinkingPart({ data }: { data: ThinkingData }) {
+  return (
+    <div className="flex items-center gap-2 text-xs text-kumo-subtle py-1">
+      <BrainIcon size={14} className="animate-pulse-dot" />
+      <span>Thinking with {data.model}&hellip;</span>
+    </div>
+  );
+}
+
+function UsagePart({ data }: { data: UsageData }) {
+  const totalTokens = data.inputTokens + data.outputTokens;
+  const latencySec = (data.latencyMs / 1000).toFixed(1);
+
+  return (
+    <div className="flex items-center gap-3 text-[11px] text-kumo-subtle mt-2 pt-2 border-t border-kumo-line">
+      <ChartBarIcon size={12} />
+      <span>{data.model}</span>
+      <span className="opacity-40">|</span>
+      <span>{totalTokens} tokens</span>
+      <span className="opacity-40">|</span>
+      <span>{latencySec}s</span>
+    </div>
+  );
+}
+
+// ── Message helpers ─────────────────────────────────────────────────
+
+/** Extract plain text from a message's parts. */
+function getMessageText(message: ChatMessage): string {
   return message.parts
     .filter((part) => part.type === "text")
-    .map((part) => (part as { type: "text"; text: string }).text)
+    .map((part) => part.text)
     .join("");
 }
 
@@ -40,6 +141,11 @@ function Chat() {
   const [input, setInput] = useState("");
   const messagesEndRef = useRef<HTMLDivElement>(null);
 
+  // Transient data parts are not added to message.parts, they only
+  // fire the onData callback. We store the latest thinking part in
+  // local state so we can render it while streaming.
+  const [thinkingData, setThinkingData] = useState<ThinkingData | null>(null);
+
   const handleOpen = useCallback(() => setConnectionStatus("connected"), []);
   const handleClose = useCallback(
     () => setConnectionStatus("disconnected"),
@@ -58,13 +164,31 @@ function Chat() {
     onError: handleError
   });
 
-  const { messages, sendMessage, clearHistory, status } = useAgentChat({
-    agent
+  const { messages, sendMessage, clearHistory, status } = useAgentChat<
+    unknown,
+    ChatMessage
+  >({
+    agent,
+    onData(part) {
+      // Capture transient thinking parts from the onData callback.
+      // These are ephemeral — not persisted and not in message.parts.
+      if (part.type === "data-thinking") {
+        // part.data is typed as ThinkingData here — no cast needed
+        setThinkingData(part.data);
+      }
+    }
   });
 
   const isStreaming = status === "streaming";
   const isConnected = connectionStatus === "connected";
 
+  // Clear transient thinking state when streaming ends
+  useEffect(() => {
+    if (!isStreaming) {
+      setThinkingData(null);
+    }
+  }, [isStreaming]);
+
   useEffect(() => {
     messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
   }, [messages]);
@@ -139,15 +263,40 @@ function Chat() {
               );
             }
 
+            // Transient parts (like data-thinking) are not in message.parts,
+            // they're captured via onData and stored in local state instead.
+            const sourcesPart = message.parts.find(
+              (p) => p.type === "data-sources"
+            );
+            const usagePart = message.parts.find(
+              (p) => p.type === "data-usage"
+            );
+
             return (
               <div key={message.id} className="flex justify-start">
                 <div className="max-w-[85%] px-4 py-2.5 rounded-2xl rounded-bl-md bg-kumo-base text-kumo-default leading-relaxed">
+                  {sourcesPart && (
+                    <SourcesPart
+                      data={sourcesPart.data}
+                      isStreaming={isLastAssistant && isStreaming}
+                    />
+                  )}
+
+                  {/* Transient thinking indicator that is captured via onData and
+                      only visible on the last assistant message while streaming */}
+                  {thinkingData && isLastAssistant && isStreaming && (
+                    <ThinkingPart data={thinkingData} />
+                  )}
+
+                  {/* Message text */}
                   <div className="whitespace-pre-wrap">
                     {text}
                     {isLastAssistant && isStreaming && (
                       <span className="inline-block w-0.5 h-[1em] bg-kumo-brand ml-0.5 align-text-bottom animate-blink-cursor" />
                     )}
                   </div>
+
+                  {usagePart && <UsagePart data={usagePart.data} />}
                 </div>
               </div>
             );