fix(think): preserve regeneration branches after restart

packages/agents/src/chat/resumable-stream.ts

@@ -90,7 +90,9 @@ function unpackSegmentBody(rowBody: string): string[] {
 function isMissingMetadataColumnError(error: unknown): boolean {
   const message = error instanceof Error ? error.message : String(error);
   return (
-    (message.includes("message_id") || message.includes("is_continuation")) &&
+    (message.includes("message_id") ||
+      message.includes("is_continuation") ||
+      message.includes("parent_message_id")) &&
     (message.toLowerCase().includes("no such column") ||
       message.toLowerCase().includes("has no column named"))
   );
@@ -124,6 +126,11 @@ type StreamMetadata = {
    * legacy rows written before this column existed.
    */
   message_id: string | null;
+  /**
+   * Parent message for a newly reconstructed assistant message. Null for flat
+   * histories, linear appends, and rows written before branch-aware recovery.
+   */
+  parent_message_id: string | null;
   /**
    * Whether this stream is a continuation (appends to the last assistant
    * message rather than starting a new one). Live broadcast frames carry
@@ -190,6 +197,7 @@ export class ResumableStream {
       created_at integer not null,
       completed_at integer,
       message_id text,
+      parent_message_id text,
       is_continuation integer
     )`;
 
@@ -217,6 +225,13 @@ export class ResumableStream {
       this
         .sql`alter table cf_ai_chat_stream_metadata add column message_id text`;
     }
+    const hasParentMessageId = columns.some(
+      (column) => column.name === "parent_message_id"
+    );
+    if (!hasParentMessageId) {
+      this
+        .sql`alter table cf_ai_chat_stream_metadata add column parent_message_id text`;
+    }
     const hasIsContinuation = columns.some(
       (column) => column.name === "is_continuation"
     );
@@ -258,7 +273,11 @@ export class ResumableStream {
    */
   start(
     requestId: string,
-    options: { messageId?: string; continuation?: boolean } = {}
+    options: {
+      messageId?: string;
+      parentMessageId?: string;
+      continuation?: boolean;
+    } = {}
   ): string {
     // Flush any pending chunks from previous streams to prevent mixing
     this.flushBuffer();
@@ -271,19 +290,30 @@ export class ResumableStream {
     this._activeIsContinuation = options.continuation ?? false;
 
     const messageId = options.messageId ?? null;
-
-    try {
+    const parentMessageId = options.parentMessageId ?? null;
+
+    const insertMetadata = () => {
+      if (parentMessageId === null) {
+        // Flat histories and linear Think turns do not need the new column, so
+        // existing databases migrate only when a branch-scoped stream needs it.
+        this.sql`
+          insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at, message_id, is_continuation)
+          values (${streamId}, ${requestId}, 'streaming', ${Date.now()}, ${messageId}, ${this._activeIsContinuation ? 1 : 0})
+        `;
+        return;
+      }
       this.sql`
-        insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at, message_id, is_continuation)
-        values (${streamId}, ${requestId}, 'streaming', ${Date.now()}, ${messageId}, ${this._activeIsContinuation ? 1 : 0})
+        insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at, message_id, parent_message_id, is_continuation)
+        values (${streamId}, ${requestId}, 'streaming', ${Date.now()}, ${messageId}, ${parentMessageId}, ${this._activeIsContinuation ? 1 : 0})
       `;
+    };
+
+    try {
+      insertMetadata();
     } catch (error) {
       if (!isMissingMetadataColumnError(error)) throw error;
       this._migrateMetadataColumns();
-      this.sql`
-        insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at, message_id, is_continuation)
-        values (${streamId}, ${requestId}, 'streaming', ${Date.now()}, ${messageId}, ${this._activeIsContinuation ? 1 : 0})
-      `;
+      insertMetadata();
     }
 
     return streamId;
@@ -310,6 +340,25 @@ export class ResumableStream {
     return rows[0].message_id ?? null;
   }
 
+  /**
+   * Parent message captured when the stream started. Orphan persistence uses
+   * this after restart to attach a reconstructed branch response correctly.
+   */
+  getStreamParentMessageId(streamId: string): string | null {
+    let rows: Array<{ parent_message_id: string | null }>;
+    try {
+      rows = this.sql<{ parent_message_id: string | null }>`
+        select parent_message_id from cf_ai_chat_stream_metadata
+        where id = ${streamId}
+      `;
+    } catch (error) {
+      if (!isMissingMetadataColumnError(error)) throw error;
+      return null;
+    }
+    if (!rows || rows.length === 0) return null;
+    return rows[0].parent_message_id ?? null;
+  }
+
   /**
    * Mark a stream as completed and flush any pending chunks.
    * @param streamId - The stream to mark as completed

packages/agents/src/tests/agents/session.ts

@@ -196,9 +196,10 @@ export class TestSessionAgent extends Agent {
   // ── ResumableStream legacy migration helpers ────────────────────
 
   /**
-   * Recreate `cf_ai_chat_stream_metadata` with the pre-#1691/#1733 schema
-   * (no `message_id` / `is_continuation`) and seed one in-flight row, so a
-   * fresh `ResumableStream` exercises the legacy lazy-migration path on a
+   * Recreate `cf_ai_chat_stream_metadata` with the pre-branch-parent schema
+   * (no `message_id` / `parent_message_id` / `is_continuation`) and seed one
+   * in-flight row, so a fresh `ResumableStream` exercises the legacy
+   * lazy-migration path on a
    * real workerd SQLite (validates the runtime's actual error strings).
    */
   async setupLegacyStreamTableForTest(): Promise<void> {
@@ -216,12 +217,55 @@ export class TestSessionAgent extends Agent {
       values ('legacy-stream', 'legacy-req', 'streaming', ${Date.now()})`;
   }
 
+  /**
+   * Recreate the metadata table as deployed immediately before branch-parent
+   * persistence: existing stream fields are present, but `parent_message_id`
+   * is not. A normal linear stream should continue using this schema unchanged.
+   */
+  async setupPreBranchParentStreamTableForTest(): Promise<void> {
+    this.sql`drop table if exists cf_ai_chat_stream_metadata`;
+    this.sql`drop table if exists cf_ai_chat_stream_chunks`;
+    this.sql`create table cf_ai_chat_stream_metadata (
+      id text primary key,
+      request_id text not null,
+      status text not null,
+      created_at integer not null,
+      completed_at integer,
+      message_id text,
+      is_continuation integer
+    )`;
+  }
+
   private streamMetadataColumnsForTest(): string[] {
     return this.sql<{ name: string }>`
       select name from pragma_table_info('cf_ai_chat_stream_metadata')
     `.map((c) => c.name);
   }
 
+  /** A linear start must not migrate a branch-only metadata column. */
+  async resumableLinearStartWithoutParentForTest(): Promise<{
+    startThrew: boolean;
+    columnsAfter: string[];
+  }> {
+    const stream = new ResumableStream(
+      <T = Record<string, unknown>>(
+        strings: TemplateStringsArray,
+        ...values: (string | number | boolean | null)[]
+      ): T[] => this.sql<T>(strings, ...values)
+    );
+
+    let startThrew = false;
+    try {
+      stream.start("linear-request", { messageId: "linear-message" });
+    } catch {
+      startThrew = true;
+    }
+    return {
+      startThrew,
+      columnsAfter: this.streamMetadataColumnsForTest()
+    };
+  }
+
   /**
    * Drive a `ResumableStream` over a legacy metadata table and report what
    * happened, so the test can assert the lazy migration recovered instead of
@@ -231,9 +275,11 @@ export class TestSessionAgent extends Agent {
   async resumableLegacyMigrationForTest(): Promise<{
     columnsBefore: string[];
     legacyMessageId: string | null;
+    legacyParentMessageId: string | null;
     startThrew: boolean;
     columnsAfter: string[];
     newStreamMessageId: string | null;
+    newStreamParentMessageId: string | null;
   }> {
     const stream = new ResumableStream(
       <T = Record<string, unknown>>(
@@ -245,13 +291,16 @@ export class TestSessionAgent extends Agent {
     const columnsBefore = this.streamMetadataColumnsForTest();
     // SELECT of the new column on a legacy row: guarded → null, no throw.
     const legacyMessageId = stream.getStreamMessageId("legacy-stream");
+    const legacyParentMessageId =
+      stream.getStreamParentMessageId("legacy-stream");
 
     // INSERT naming the new columns on a legacy table: must migrate + retry.
     let startThrew = false;
     let newStreamId = "";
     try {
       newStreamId = stream.start("req-x", {
         messageId: "msg-1",
+        parentMessageId: "parent-1",
         continuation: true
       });
     } catch {
@@ -262,13 +311,18 @@ export class TestSessionAgent extends Agent {
     const newStreamMessageId = newStreamId
       ? stream.getStreamMessageId(newStreamId)
       : null;
+    const newStreamParentMessageId = newStreamId
+      ? stream.getStreamParentMessageId(newStreamId)
+      : null;
 
     return {
       columnsBefore,
       legacyMessageId,
+      legacyParentMessageId,
       startThrew,
       columnsAfter,
-      newStreamMessageId
+      newStreamMessageId,
+      newStreamParentMessageId
     };
   }
 }

packages/agents/src/tests/resumable-stream-migration.test.ts

@@ -5,7 +5,8 @@ import { getAgentByName } from "..";
 /**
  * ResumableStream lazy metadata-column migration (#1691, #1733).
  *
- * New tables are created with `message_id` / `is_continuation` up front, so
+ * New tables are created with `message_id` / `parent_message_id` /
+ * `is_continuation` up front, so
  * most wakes never migrate. Tables created by an older release lack those
  * columns and must migrate lazily — on the first stream write that needs
  * them — instead of paying a schema-introspection read on every construction.
@@ -17,12 +18,19 @@ import { getAgentByName } from "..";
 
 interface LegacyMigrationStub {
   setupLegacyStreamTableForTest(): Promise<void>;
+  setupPreBranchParentStreamTableForTest(): Promise<void>;
+  resumableLinearStartWithoutParentForTest(): Promise<{
+    startThrew: boolean;
+    columnsAfter: string[];
+  }>;
   resumableLegacyMigrationForTest(): Promise<{
     columnsBefore: string[];
     legacyMessageId: string | null;
+    legacyParentMessageId: string | null;
     startThrew: boolean;
     columnsAfter: string[];
     newStreamMessageId: string | null;
+    newStreamParentMessageId: string | null;
   }>;
 }
 
@@ -39,25 +47,39 @@ describe("ResumableStream — legacy metadata-column migration", () => {
     name = `rs-migrate-${Date.now()}-${Math.random().toString(36).slice(2)}`;
   });
 
-  it("migrates a legacy table on first stream write instead of throwing", async () => {
+  it("does not migrate the branch parent column for a linear stream", async () => {
+    const agent = await getAgent(name);
+    await agent.setupPreBranchParentStreamTableForTest();
+
+    const result = await agent.resumableLinearStartWithoutParentForTest();
+
+    expect(result.startThrew).toBe(false);
+    expect(result.columnsAfter).not.toContain("parent_message_id");
+  });
+
+  it("migrates a legacy table on first branch stream write instead of throwing", async () => {
     const agent = await getAgent(name);
     await agent.setupLegacyStreamTableForTest();
 
     const result = await agent.resumableLegacyMigrationForTest();
 
     // Precondition: the seeded table really is the old schema.
     expect(result.columnsBefore).not.toContain("message_id");
+    expect(result.columnsBefore).not.toContain("parent_message_id");
     expect(result.columnsBefore).not.toContain("is_continuation");
 
-    // Reading the new column off a legacy row is guarded → null, no throw.
+    // Reading the new columns off a legacy row is guarded → null, no throw.
     expect(result.legacyMessageId).toBeNull();
+    expect(result.legacyParentMessageId).toBeNull();
 
     // start() hit the missing columns, migrated, and retried successfully.
     expect(result.startThrew).toBe(false);
     expect(result.columnsAfter).toContain("message_id");
+    expect(result.columnsAfter).toContain("parent_message_id");
     expect(result.columnsAfter).toContain("is_continuation");
 
-    // The migrated row round-trips the value start() wrote.
+    // The migrated row round-trips the values start() wrote.
     expect(result.newStreamMessageId).toBe("msg-1");
+    expect(result.newStreamParentMessageId).toBe("parent-1");
   });
 });

packages/think/src/tests/agents/think-session.ts

@@ -8089,6 +8089,57 @@ export class ThinkNonRecoveryTestAgent extends Think {
     return this.getMessages();
   }
 
+  /**
+   * Seed the durable state reached when a branch-scoped regeneration loses its
+   * live stream reader. Chat recovery is disabled on this agent, so only the
+   * client resume handshake can materialize the buffered partial after restart.
+   */
+  async seedInterruptedRegenerationForResumeTest(): Promise<{
+    requestId: string;
+    userId: string;
+    oldAssistantId: string;
+    partialAssistantId: string;
+  }> {
+    const requestId = "req-resume-regeneration";
+    const userId = "user-resume-regeneration";
+    const oldAssistantId = "assistant-old-resume-regeneration";
+    const partialAssistantId = "assistant-partial-resume-regeneration";
+
+    await this.session.appendMessage({
+      id: userId,
+      role: "user",
+      parts: [{ type: "text", text: "answer this differently" }]
+    });
+    await this.session.appendMessage({
+      id: oldAssistantId,
+      role: "assistant",
+      parts: [{ type: "text", text: "Old answer" }]
+    });
+
+    const streamId = this._startResumableStream(requestId, {
+      parentMessageId: userId
+    });
+    for (const body of [
+      JSON.stringify({ type: "start", messageId: partialAssistantId }),
+      JSON.stringify({ type: "text-start", id: "regenerated-text" }),
+      JSON.stringify({
+        type: "text-delta",
+        id: "regenerated-text",
+        delta: "Partial replacement"
+      })
+    ]) {
+      this._resumableStream.storeChunk(streamId, body);
+    }
+    this._resumableStream.flushBuffer();
+
+    return { requestId, userId, oldAssistantId, partialAssistantId };
+  }
+
+  /** Stored child branches for resume-handshake recovery assertions. */
+  async getBranchesForTest(messageId: string): Promise<UIMessage[]> {
+    return (await this.session.getBranches(messageId)) as UIMessage[];
+  }
+
   async getActiveFibers(): Promise<Array<{ id: string; name: string }>> {
     return this.sql<{ id: string; name: string }>`
       SELECT id, name FROM cf_agents_runs