fix(ai-chat): fix regenerate stale messages, transport trigger, and identical-content reconciliation

.changeset/fix-regenerate-stale-messages.md

@@ -0,0 +1,51 @@
+---
+"@cloudflare/ai-chat": patch
+---
+
+Fix `regenerate()` leaving stale assistant messages in SQLite
+
+**Bug 1 — Transport drops `trigger` field:**
+`WebSocketChatTransport.sendMessages` was not including the `trigger` field
+(e.g. `"regenerate-message"`, `"submit-message"`) in the body payload sent
+to the server. The AI SDK passes this field so the server can distinguish
+between a new message and a regeneration request. Fixed by adding
+`trigger: options.trigger` to the serialized body.
+
+On the server side, `trigger` is now destructured out of the parsed body
+alongside `messages` and `clientTools`, so it does not leak into
+`options.body` in `onChatMessage`. Users who inspect `options.body` will
+not see any change in behavior.
+
+**Bug 2 — `persistMessages` never deletes stale rows:**
+`persistMessages` only performed `INSERT ... ON CONFLICT DO UPDATE` (upsert),
+so when `regenerate()` removed the last assistant message from the client's
+array, the old row persisted in SQLite. On the next `_loadMessagesFromDb`,
+the stale assistant message reappeared in `this.messages`, causing:
+
+- Anthropic models to reject with HTTP 400 (conversation must end with a
+  user message)
+- Duplicate/phantom assistant messages across reconnects
+
+Fixed by adding an internal `_deleteStaleRows` option to `persistMessages`.
+When the chat-request handler (`CF_AGENT_USE_CHAT_REQUEST`) calls
+`persistMessages`, it passes `{ _deleteStaleRows: true }`, which deletes
+any DB rows whose IDs are absent from the incoming (post-merge) message set.
+This uses the post-merge IDs from `_mergeIncomingWithServerState` to
+correctly handle cases where client assistant IDs are remapped to server IDs.
+
+The `_deleteStaleRows` flag is internal only (`@internal` JSDoc) and is
+never passed by user code or other handlers (`CF_AGENT_CHAT_MESSAGES`,
+`_reply`, `saveMessages`). The default behavior of `persistMessages`
+(upsert-only, no deletes) is unchanged.
+
+**Bug 3 — Content-based reconciliation mismatches identical messages:**
+`_reconcileAssistantIdsWithServerState` used a single-pass cursor for both
+exact-ID and content-based matching. When an exact-ID match jumped the
+cursor forward, it skipped server messages needed for content matching
+of later identical-text assistant messages (e.g. "Sure", "I understand").
+
+Rewritten with a two-pass approach: Pass 1 resolves all exact-ID matches
+and claims server indices. Pass 2 does content-based matching only over
+unclaimed server indices. This prevents exact-ID matches from interfering
+with content matching, fixing duplicate rows in long conversations with
+repeated short assistant responses.

packages/ai-chat/src/index.ts

@@ -419,9 +419,15 @@ export class AIChatAgent<
             return;
           }
 
-          const { messages, clientTools, ...customBody } = parsed as {
+          const {
+            messages,
+            clientTools,
+            trigger: _trigger,
+            ...customBody
+          } = parsed as {
             messages: ChatMessage[];
             clientTools?: ClientToolSchema[];
+            trigger?: string;
             [key: string]: unknown;
           };
 
@@ -442,7 +448,9 @@ export class AIChatAgent<
             [connection.id]
           );
 
-          await this.persistMessages(transformedMessages, [connection.id]);
+          await this.persistMessages(transformedMessages, [connection.id], {
+            _deleteStaleRows: true
+          });
 
           this.observability?.emit(
             {
@@ -1043,7 +1051,9 @@ export class AIChatAgent<
 
   async persistMessages(
     messages: ChatMessage[],
-    excludeBroadcastIds: string[] = []
+    excludeBroadcastIds: string[] = [],
+    /** @internal */
+    options?: { _deleteStaleRows?: boolean }
   ) {
     // Merge incoming messages with existing server state to preserve tool outputs.
     // This is critical for client-side tools: the client sends messages without
@@ -1071,6 +1081,30 @@ export class AIChatAgent<
       this._persistedMessageCache.set(safe.id, json);
     }
 
+    // Reconcile: delete DB rows not present in the incoming message set.
+    // The transport always sends the full message array, so any DB row
+    // absent from it is stale (e.g. regenerate() removes the last assistant
+    // message). Without this, the stale row stays in SQLite and gets
+    // reloaded into this.messages, causing providers like Anthropic to
+    // reject with 400 (conversation must end with a user message).
+    // This MUST use mergedMessages (post-merge IDs) because
+    // _mergeIncomingWithServerState can remap client IDs to server IDs.
+    if (options?._deleteStaleRows) {
+      const keepIds = new Set(mergedMessages.map((m) => m.id));
+      const allDbRows =
+        this.sql<{ id: string }>`
+          select id from cf_ai_chat_agent_messages
+        ` || [];
+      for (const row of allDbRows) {
+        if (!keepIds.has(row.id)) {
+          this.sql`
+            delete from cf_ai_chat_agent_messages where id = ${row.id}
+          `;
+          this._persistedMessageCache.delete(row.id);
+        }
+      }
+    }
+
     // Enforce maxPersistedMessages: delete oldest messages if over the limit
     if (this.maxPersistedMessages != null) {
       this._enforceMaxPersistedMessages();
@@ -1164,8 +1198,18 @@ export class AIChatAgent<
    * The client can keep a different local ID for an assistant message than the one
    * persisted on the server (e.g. optimistic/local IDs). When that full history is
    * sent back, persisting by ID alone creates duplicate assistant rows. To prevent
-   * this, we reuse the server ID for assistant messages that match by content and
-   * order, while leaving tool-call messages to _resolveMessageForToolMerge.
+   * this, we reuse the server ID for assistant messages that match by content,
+   * while leaving tool-call messages to _resolveMessageForToolMerge.
+   *
+   * Uses a two-pass approach:
+   *  - Pass 1: resolve all exact-ID matches, claiming server indices.
+   *  - Pass 2: content-based matching for remaining non-tool assistant messages,
+   *    scanning only unclaimed server indices left-to-right.
+   *
+   * The two-pass design prevents exact-ID matches from advancing a cursor past
+   * server messages that a later incoming message needs for content matching.
+   * This fixes mismatches when two assistant messages have identical text
+   * (e.g. "Sure", "I understand") — see #1008.
    */
   private _reconcileAssistantIdsWithServerState(
     incomingMessages: ChatMessage[]
@@ -1174,22 +1218,27 @@ export class AIChatAgent<
       return incomingMessages;
     }
 
-    // Tracks the earliest server index we should consider for subsequent matches.
-    // This preserves ordering and prevents one server message from being reused for
-    // multiple incoming messages with identical content.
-    let serverCursor = 0;
-
-    return incomingMessages.map((incomingMessage) => {
-      // Fast path: exact ID already exists in server history.
-      // This applies to any role (user/assistant/system/tool), so in-order
-      // round-trips naturally advance the cursor even when assistant content
-      // reconciliation is skipped.
-      const exactMatchIndex = this.messages.findIndex(
-        (serverMessage, index) =>
-          index >= serverCursor && serverMessage.id === incomingMessage.id
+    // Pass 1: Resolve exact-ID matches first.
+    // This prevents content-based matching from claiming a server message
+    // that has a direct ID match with a later incoming message.
+    const claimedServerIndices = new Set<number>();
+    const exactMatchMap = new Map<number, number>();
+
+    for (let i = 0; i < incomingMessages.length; i++) {
+      const serverIdx = this.messages.findIndex(
+        (sm, si) =>
+          !claimedServerIndices.has(si) && sm.id === incomingMessages[i].id
       );
-      if (exactMatchIndex !== -1) {
-        serverCursor = exactMatchIndex + 1;
+      if (serverIdx !== -1) {
+        claimedServerIndices.add(serverIdx);
+        exactMatchMap.set(i, serverIdx);
+      }
+    }
+
+    // Pass 2: Content-based matching for remaining non-tool assistant messages.
+    // Scans unclaimed server messages left-to-right to preserve ordering.
+    return incomingMessages.map((incomingMessage, incomingIdx) => {
+      if (exactMatchMap.has(incomingIdx)) {
         return incomingMessage;
       }
 
@@ -1207,7 +1256,10 @@ export class AIChatAgent<
         return incomingMessage;
       }
 
-      for (let i = serverCursor; i < this.messages.length; i++) {
+      for (let i = 0; i < this.messages.length; i++) {
+        if (claimedServerIndices.has(i)) {
+          continue;
+        }
         const serverMessage = this.messages[i];
         if (
           serverMessage.role !== "assistant" ||
@@ -1217,7 +1269,7 @@ export class AIChatAgent<
         }
 
         if (this._assistantMessageContentKey(serverMessage) === incomingKey) {
-          serverCursor = i + 1;
+          claimedServerIndices.add(i);
 
           return {
             ...incomingMessage,

packages/ai-chat/src/tests/chat-persistence.test.ts

@@ -55,6 +55,17 @@ describe("Chat Agent Persistence", () => {
     const firstDone = await donePromise;
     expect(firstDone).toBe(true);
 
+    // Fetch persisted messages to capture the assistant response from the
+    // first request. In a real AI SDK flow, the client always sends the full
+    // message array including previous assistant messages.
+    const midReq = new Request(
+      `http://example.com/agents/test-chat-agent/${room}/get-messages`
+    );
+    const midRes = await worker.fetch(midReq, env, createExecutionContext());
+    const midMessages = (await midRes.json()) as ChatMessage[];
+    const firstAssistant = midMessages.find((m) => m.role === "assistant");
+    expect(firstAssistant).toBeDefined();
+
     const secondMessage: ChatMessage = {
       id: "msg2",
       role: "user",
@@ -74,13 +85,17 @@ describe("Chat Agent Persistence", () => {
       }
     });
 
+    // Include the first assistant message in the second request (mirrors
+    // real AI SDK behavior — the client always sends all messages).
     ws.send(
       JSON.stringify({
         type: MessageType.CF_AGENT_USE_CHAT_REQUEST,
         id: "req2",
         init: {
           method: "POST",
-          body: JSON.stringify({ messages: [firstMessage, secondMessage] })
+          body: JSON.stringify({
+            messages: [firstMessage, firstAssistant!, secondMessage]
+          })
         }
       })
     );