ai-chat: validate persisted messages, fix metadata on broadcast/resume, document advanced patterns

.changeset/validate-messages-and-metadata.md

@@ -0,0 +1,13 @@
+---
+"@cloudflare/ai-chat": patch
+---
+
+Add structural message validation and fix message metadata on broadcast/resume path.
+
+**Structural message validation:**
+
+Messages loaded from SQLite are now validated for required structure (non-empty `id` string, valid `role`, `parts` is an array). Malformed rows — from corruption, manual tampering, or schema drift — are logged with a warning and silently skipped instead of crashing the agent. This is intentionally lenient: empty `parts` arrays are allowed (streams that errored mid-flight), and no tool/data schema validation is performed at load time (that remains a userland concern via `safeValidateUIMessages` from the AI SDK).
+
+**Message metadata on broadcast/resume path:**
+
+The server already captures `messageMetadata` from `start`, `finish`, and `message-metadata` stream chunks and persists it on `message.metadata`. However, the client-side broadcast path (multi-tab sync) and stream resume path (reconnection) did not propagate metadata — the `activeStreamRef` only tracked `parts`. Now it also tracks `metadata`, and `flushActiveStreamToMessages` includes it in the partial message flushed to React state. This means cross-tab clients and reconnecting clients see metadata (model info, token usage, timestamps) during streaming, not just after the final `CF_AGENT_CHAT_MESSAGES` broadcast.

docs/chat-agents.md

@@ -614,6 +614,220 @@ const result = streamText({
 });
 ```
 
+## Advanced Patterns
+
+Since `onChatMessage` gives you full control over the `streamText` call, you can use any AI SDK feature directly. The patterns below all work out of the box — no special `AIChatAgent` configuration is needed.
+
+### Dynamic Model and Tool Control
+
+Use [`prepareStep`](https://ai-sdk.dev/docs/agents/loop-control) to change the model, available tools, or system prompt between steps in a multi-step agent loop:
+
+```typescript
+import { streamText, convertToModelMessages, tool, stepCountIs } from "ai";
+import { z } from "zod";
+
+async onChatMessage() {
+  const result = streamText({
+    model: cheapModel, // Default model for simple steps
+    messages: await convertToModelMessages(this.messages),
+    tools: {
+      search: searchTool,
+      analyze: analyzeTool,
+      summarize: summarizeTool
+    },
+    stopWhen: stepCountIs(10),
+    prepareStep: async ({ stepNumber, messages }) => {
+      // Phase 1: Search (steps 0-2)
+      if (stepNumber <= 2) {
+        return {
+          activeTools: ["search"],
+          toolChoice: "required" // Force tool use
+        };
+      }
+
+      // Phase 2: Analyze with a stronger model (steps 3-5)
+      if (stepNumber <= 5) {
+        return {
+          model: expensiveModel,
+          activeTools: ["analyze"]
+        };
+      }
+
+      // Phase 3: Summarize
+      return { activeTools: ["summarize"] };
+    }
+  });
+
+  return result.toUIMessageStreamResponse();
+}
+```
+
+`prepareStep` runs before each step and can return overrides for `model`, `activeTools`, `toolChoice`, `system`, and `messages`. Use it to:
+
+- **Switch models** — use a cheap model for simple steps, escalate for reasoning
+- **Phase tools** — restrict which tools are available at each step
+- **Manage context** — prune or transform messages to stay within token limits
+- **Force tool calls** — use `toolChoice: { type: "tool", toolName: "search" }` to require a specific tool
+
+### Language Model Middleware
+
+Use [`wrapLanguageModel`](https://ai-sdk.dev/docs/ai-sdk-core/middleware) to add guardrails, RAG, caching, or logging without modifying your chat logic:
+
+```typescript
+import { streamText, convertToModelMessages, wrapLanguageModel } from "ai";
+import type { LanguageModelV3Middleware } from "@ai-sdk/provider";
+
+const guardrailMiddleware: LanguageModelV3Middleware = {
+  wrapGenerate: async ({ doGenerate }) => {
+    const { text, ...rest } = await doGenerate();
+    // Filter PII or sensitive content from the response
+    const cleaned = text?.replace(/\b\d{3}-\d{2}-\d{4}\b/g, "[REDACTED]");
+    return { text: cleaned, ...rest };
+  }
+};
+
+async onChatMessage() {
+  const model = wrapLanguageModel({
+    model: baseModel,
+    middleware: [guardrailMiddleware]
+  });
+
+  const result = streamText({
+    model,
+    messages: await convertToModelMessages(this.messages)
+  });
+
+  return result.toUIMessageStreamResponse();
+}
+```
+
+The AI SDK includes built-in middlewares:
+
+- `extractReasoningMiddleware` — surface chain-of-thought from models like DeepSeek R1
+- `defaultSettingsMiddleware` — apply default temperature, max tokens, etc.
+- `simulateStreamingMiddleware` — add streaming to non-streaming models
+
+Multiple middlewares compose in order: `middleware: [first, second]` applies as `first(second(model))`.
+
+### Structured Output
+
+Use [`generateObject`](https://ai-sdk.dev/docs/ai-sdk-core/generating-structured-data) inside tools for structured data extraction:
+
+```typescript
+import {
+  streamText, generateObject, convertToModelMessages, tool, stepCountIs
+} from "ai";
+import { z } from "zod";
+
+async onChatMessage() {
+  const result = streamText({
+    model: myModel,
+    messages: await convertToModelMessages(this.messages),
+    tools: {
+      extractContactInfo: tool({
+        description: "Extract structured contact information from the conversation",
+        inputSchema: z.object({
+          text: z.string().describe("The text to extract contact info from")
+        }),
+        execute: async ({ text }) => {
+          const { object } = await generateObject({
+            model: myModel,
+            schema: z.object({
+              name: z.string(),
+              email: z.string().email(),
+              phone: z.string().optional()
+            }),
+            prompt: `Extract contact information from: ${text}`
+          });
+          return object;
+        }
+      })
+    },
+    stopWhen: stepCountIs(5)
+  });
+
+  return result.toUIMessageStreamResponse();
+}
+```
+
+### Subagent Delegation
+
+Tools can delegate work to focused sub-calls with their own context. Use [`ToolLoopAgent`](https://ai-sdk.dev/docs/reference/ai-sdk-core/tool-loop-agent) to define a reusable agent, then call it from a tool's `execute`:
+
+```typescript
+import {
+  ToolLoopAgent, streamText, convertToModelMessages, tool, stepCountIs
+} from "ai";
+import { z } from "zod";
+
+// Define a reusable research agent with its own tools and instructions
+const researchAgent = new ToolLoopAgent({
+  model: researchModel,
+  instructions: "You are a research assistant. Be thorough and cite sources.",
+  tools: { webSearch: webSearchTool },
+  stopWhen: stepCountIs(10)
+});
+
+async onChatMessage() {
+  const result = streamText({
+    model: orchestratorModel,
+    messages: await convertToModelMessages(this.messages),
+    tools: {
+      deepResearch: tool({
+        description: "Research a topic in depth",
+        inputSchema: z.object({
+          topic: z.string().describe("The topic to research")
+        }),
+        execute: async ({ topic }) => {
+          const { text } = await researchAgent.generate({ prompt: topic });
+          return { summary: text };
+        }
+      })
+    },
+    stopWhen: stepCountIs(5)
+  });
+
+  return result.toUIMessageStreamResponse();
+}
+```
+
+The research agent runs in its own context — its token budget is separate from the orchestrator's. Only the summary goes back to the parent model.
+
+`ToolLoopAgent` is best suited for subagents, not as a replacement for `streamText` in `onChatMessage` itself. The main `onChatMessage` benefits from direct access to `this.env`, `this.messages`, and `options.body` — things that a pre-configured `ToolLoopAgent` instance cannot reference.
+
+#### Streaming progress with preliminary results
+
+By default, a tool part appears as loading until `execute` returns. Use an async generator (`async function*`) to stream progress updates to the client while the tool is still working:
+
+```typescript
+deepResearch: tool({
+  description: "Research a topic in depth",
+  inputSchema: z.object({
+    topic: z.string().describe("The topic to research")
+  }),
+  async *execute({ topic }) {
+    // Preliminary result — the client sees "searching" immediately
+    yield { status: "searching", topic, summary: undefined };
+
+    const { text } = await researchAgent.generate({ prompt: topic });
+
+    // Final result — sent to the model for its next step
+    yield { status: "done", topic, summary: text };
+  }
+});
+```
+
+Each `yield` updates the tool part on the client in real-time (with `preliminary: true`). The last yielded value becomes the final output that the model sees. This is useful for long-running tools where you want to show status (searching, analyzing, summarizing) as the work progresses.
+
+This pattern is useful when:
+
+- A task requires exploring large amounts of information that would bloat the main context
+- You want to show real-time progress for long-running tools
+- You want to parallelize independent research (multiple tool calls run concurrently)
+- You need different models or system prompts for different subtasks
+
+For more, see the [AI SDK Agents docs](https://ai-sdk.dev/docs/agents/overview), [Subagents](https://ai-sdk.dev/docs/agents/subagents), and [Preliminary Tool Results](https://ai-sdk.dev/docs/ai-sdk-core/tools-and-tool-calling#preliminary-tool-results).
+
 ## Multi-Client Sync
 
 When multiple clients connect to the same agent instance, messages are automatically broadcast to all connections. If one client sends a message, all other connected clients receive the updated message list.

packages/ai-chat/src/index.ts

@@ -32,6 +32,32 @@ import { nanoid } from "nanoid";
 /** Shared encoder for UTF-8 byte length measurement */
 const textEncoder = new TextEncoder();
 
+/**
+ * Validates that a parsed message has the minimum required structure.
+ * Returns false for messages that would cause runtime errors downstream
+ * (e.g. in convertToModelMessages or the UI layer).
+ *
+ * Checks:
+ * - `id` is a non-empty string
+ * - `role` is one of the valid roles
+ * - `parts` is an array (may be empty — the AI SDK enforces nonempty
+ *   on incoming messages, but we are lenient on persisted data)
+ */
+function isValidMessageStructure(msg: unknown): msg is ChatMessage {
+  if (typeof msg !== "object" || msg === null) return false;
+  const m = msg as Record<string, unknown>;
+
+  if (typeof m.id !== "string" || m.id.length === 0) return false;
+
+  if (m.role !== "user" && m.role !== "assistant" && m.role !== "system") {
+    return false;
+  }
+
+  if (!Array.isArray(m.parts)) return false;
+
+  return true;
+}
+
 /**
  * One-shot deprecation warnings (warns once per key per session).
  */
@@ -265,7 +291,12 @@ export class AIChatAgent<
     // Initialize resumable stream manager (creates its own tables + restores state)
     this._resumableStream = new ResumableStream(this.sql.bind(this));
 
-    // Load messages and automatically transform them to v5 format
+    // Load messages and automatically transform them to v5 format.
+    // Note: _loadMessagesFromDb() runs structural validation which requires
+    // `parts` to be an array. Legacy v4 messages (with `content` instead of
+    // `parts`) would fail this check — but that's fine because autoTransformMessages
+    // already migrated them on a previous load, and persistMessages wrote them back.
+    // Any message still without `parts` at this point is genuinely corrupt.
     const rawMessages = this._loadMessagesFromDb();
 
     // Automatic migration following https://jhak.im/blog/ai-sdk-migration-handling-previously-saved-messages
@@ -684,6 +715,18 @@ export class AIChatAgent<
         try {
           const messageStr = row.message as string;
           const parsed = JSON.parse(messageStr) as ChatMessage;
+
+          // Structural validation: ensure required fields exist and have
+          // the correct types. This catches corrupted rows, manual tampering,
+          // or schema drift from older versions without crashing the agent.
+          if (!isValidMessageStructure(parsed)) {
+            console.warn(
+              `[AIChatAgent] Skipping invalid message ${row.id}: ` +
+                "missing or malformed id, role, or parts"
+            );
+            return null;
+          }
+
           // Cache the raw JSON keyed by message ID
           this._persistedMessageCache.set(parsed.id, messageStr);
           return parsed;