fix(think): preserve regeneration branch context

.changeset/fix-think-regeneration-context.md

@@ -0,0 +1,8 @@
+---
+"agents": patch
+"@cloudflare/think": patch
+---
+
+Generate Think response alternatives from the selected regeneration branch point instead of continuing the previous answer.
+
+Preserve the branch selection through compaction and chat recovery so interrupted regenerated responses remain siblings of the previous answer.

packages/agents/src/chat/orphan-persist.ts

@@ -39,6 +39,8 @@ export interface PersistReconstructedOrphanOptions<
    * present.
    */
   fallbackId: string;
+  /** Parent for a newly reconstructed message in tree-structured stores. */
+  parentId?: string | null;
   /**
    * Finalize the reconstructed message before upsert — e.g. strip internal
    * parts or resolve the persist-target id. Return `null` to skip persistence
@@ -84,7 +86,7 @@ export async function persistReconstructedOrphan<
   if (existing) {
     await options.store.updateMessage(options.merge(existing, prepared));
   } else {
-    await options.store.appendMessage(prepared);
+    await options.store.appendMessage(prepared, options.parentId);
   }
   return true;
 }

packages/agents/src/chat/recovery-engine.ts

@@ -352,7 +352,10 @@ export interface ChatFiberWakeHooks<TClassify> {
     input: PersistOrphanedPartialInput
   ): boolean | Promise<boolean>;
   /** Materialize the orphaned stream's partial into a persisted assistant message. */
-  persistOrphanedStream(streamId: string): Promise<void>;
+  persistOrphanedStream(
+    streamId: string,
+    snapshot: ChatFiberSnapshot | null
+  ): Promise<void>;
   /** Mark the (still-active) recovered stream complete and schedule cleanup. */
   completeRecoveredStream(streamId: string): void | Promise<void>;
   /**
@@ -469,7 +472,7 @@ export class ChatRecoveryEngine {
           partial
         })
       ) {
-        await wake.persistOrphanedStream(streamId);
+        await wake.persistOrphanedStream(streamId, snapshot);
       }
       await adapter.exhaustChatRecovery(
         incident,
@@ -508,7 +511,7 @@ export class ChatRecoveryEngine {
           partial
         })
       ) {
-        await wake.persistOrphanedStream(streamId);
+        await wake.persistOrphanedStream(streamId, snapshot);
       }
 
       if (streamStillActive) {

packages/agents/src/chat/recovery.ts

@@ -23,6 +23,10 @@ export type ChatFiberSnapshot<Kind extends string = string> = {
   latestMessageId?: string;
   latestMessageRole?: string;
   latestUserMessageId?: string;
+  /** Explicit history endpoint for a branch-scoped turn such as regeneration. */
+  historyLeafId?: string;
+  /** Physical active leaf observed when a branch-scoped turn was accepted. */
+  activeLeafIdAtStart?: string;
   startedAt: number;
   lastBody?: Record<string, unknown>;
   lastClientTools?: ClientToolSchema[];
@@ -34,6 +38,8 @@ export function createChatFiberSnapshot<Kind extends string>({
   recoveryRootRequestId,
   continuation,
   messages,
+  historyLeafId,
+  activeLeafIdAtStart,
   lastBody,
   lastClientTools
 }: {
@@ -42,6 +48,8 @@ export function createChatFiberSnapshot<Kind extends string>({
   recoveryRootRequestId?: string;
   continuation: boolean;
   messages: ReadonlyArray<SnapshotMessage>;
+  historyLeafId?: string;
+  activeLeafIdAtStart?: string;
   lastBody?: Record<string, unknown>;
   lastClientTools?: ClientToolSchema[];
 }): ChatFiberSnapshot<Kind> {
@@ -65,6 +73,8 @@ export function createChatFiberSnapshot<Kind extends string>({
     latestMessageId: latestMessage?.id,
     latestMessageRole: latestMessage?.role,
     latestUserMessageId: latestUser?.id,
+    historyLeafId,
+    activeLeafIdAtStart,
     startedAt: Date.now(),
     lastBody,
     lastClientTools

packages/agents/src/experimental/memory/session/session.ts

@@ -25,7 +25,10 @@ import {
 } from "./context";
 import { AgentSessionProvider, type SqlProvider } from "./providers/agent";
 import { AgentContextProvider } from "./providers/agent-context";
-import type { CompactResult } from "../utils/compaction-helpers";
+import {
+  COMPACTION_PREFIX,
+  type CompactResult
+} from "../utils/compaction-helpers";
 import { estimateMessageTokens, estimateStringTokens } from "../utils/tokens";
 import { MessageType } from "../../../types";
 
@@ -732,9 +735,10 @@ export class Session {
 
   /**
    * Run the registered compaction function and store the result as an overlay.
-   * Requires `onCompaction()` to be called first.
+   * When `leafId` is provided, compact that root-to-leaf branch instead of the
+   * active branch. Requires `onCompaction()` to be called first.
    */
-  async compact(): Promise<CompactResult | null> {
+  async compact(leafId?: string | null): Promise<CompactResult | null> {
     await this._ensureRestored();
     if (!this._compactionFn) {
       throw new Error(
@@ -743,14 +747,15 @@ export class Session {
     }
 
     const tokensBefore = await this._emitStatus("compacting");
+    const history = await this.getHistory(leafId);
 
     let result: CompactResult | null;
     try {
       // Pass the Session's authoritative token counter so the compaction
       // function's boundary logic can use the same accounting as the
       // fire/no-fire decision (see CompactContext). The function still wins if
       // it was given its own explicit counter.
-      result = await this._compactionFn(await this.getHistory(), {
+      result = await this._compactionFn(history, {
         tokenCounter: this._tokenCounter
       });
     } catch (err) {
@@ -763,15 +768,21 @@ export class Session {
       return null;
     }
 
-    // Validate toMessageId exists in the history
-    const historyIds = new Set((await this.getHistory()).map((m) => m.id));
+    // Validate toMessageId exists in the selected history.
+    const historyIds = new Set(history.map((message) => message.id));
     if (!historyIds.has(result.toMessageId)) {
       await this._emitStatus("idle");
       return null;
     }
 
-    // Iterative compaction — extend from earliest existing compaction's start
-    const existing = await this.getCompactions();
+    // Iterative compaction extends only an overlay visible on this branch.
+    // Other sibling branches may have unrelated overlays in the same session.
+    const existing = (await this.getCompactions()).filter(
+      (compaction) =>
+        historyIds.has(`${COMPACTION_PREFIX}${compaction.id}`) ||
+        (historyIds.has(compaction.fromMessageId) &&
+          historyIds.has(compaction.toMessageId))
+    );
     const fromId =
       existing.length > 0 ? existing[0].fromMessageId : result.fromMessageId;
 

packages/agents/src/tests/experimental/memory/session/session.test.ts

@@ -638,14 +638,19 @@ describe("ContextBlocks — edge cases", () => {
 // ── Compaction tests ─────────────────────────────────────────────
 
 function createCompactableSession(
-  compactFn: (msgs: SessionMessage[]) => Promise<CompactResult | null>
+  compactFn: (msgs: SessionMessage[]) => Promise<CompactResult | null>,
+  selectHistory?: (
+    messages: SessionMessage[],
+    leafId?: string | null
+  ) => SessionMessage[]
 ) {
   const messages: SessionMessage[] = [];
   const compactions: StoredCompaction[] = [];
 
   const storage: SessionProvider = {
     getMessage: (id) => messages.find((m) => m.id === id) ?? null,
-    getHistory: () => messages,
+    getHistory: (leafId) =>
+      selectHistory ? selectHistory(messages, leafId) : messages,
     getLatestLeaf: () => messages[messages.length - 1] ?? null,
     getBranches: () => [],
     getPathLength: () => messages.length,
@@ -736,6 +741,59 @@ describe("Session.compact()", () => {
     expect(compactions[0].toMessageId).toBe("m3");
   });
 
+  it("compacts an explicit branch without extending a sibling overlay", async () => {
+    const compactedMessageIds: string[][] = [];
+    const { session, messages, compactions } = createCompactableSession(
+      async (history): Promise<CompactResult> => {
+        compactedMessageIds.push(history.map((message) => message.id));
+        return {
+          fromMessageId: "m1",
+          toMessageId: "m1",
+          summary: "Selected branch summary"
+        };
+      },
+      (allMessages, leafId) =>
+        leafId === "m2" ? allMessages.slice(0, 3) : allMessages
+    );
+
+    messages.push(
+      { id: "m0", role: "user", parts: [{ type: "text", text: "start" }] },
+      {
+        id: "m1",
+        role: "assistant",
+        parts: [{ type: "text", text: "first answer" }]
+      },
+      {
+        id: "m2",
+        role: "user",
+        parts: [{ type: "text", text: "follow-up" }]
+      },
+      {
+        id: "old-assistant",
+        role: "assistant",
+        parts: [{ type: "text", text: "superseded answer" }]
+      }
+    );
+
+    compactions.push({
+      id: "sibling-compaction",
+      fromMessageId: "m0",
+      toMessageId: "old-assistant",
+      summary: "Sibling branch summary",
+      createdAt: new Date().toISOString()
+    });
+
+    await session.compact("m2");
+
+    expect(compactedMessageIds).toEqual([["m0", "m1", "m2"]]);
+    expect(compactions).toHaveLength(2);
+    expect(compactions[1]).toMatchObject({
+      fromMessageId: "m1",
+      toMessageId: "m1",
+      summary: "Selected branch summary"
+    });
+  });
+
   it("returns null when compaction function returns null", async () => {
     const { session, messages, compactions } = createCompactableSession(
       async () => null

packages/think/src/tests/agents/client-tools.ts

@@ -8,7 +8,8 @@
 import type { LanguageModel, ToolSet, UIMessage } from "ai";
 import { tool } from "ai";
 import { z } from "zod";
-import { Think } from "../../think";
+import { Session } from "agents/experimental/memory/session";
+import { defaultContextOverflowClassifier, Think } from "../../think";
 import type {
   ChatResponseResult,
   MessageConcurrency,
@@ -630,7 +631,10 @@ function createMultiStepExecutableClientToolMockModel(): LanguageModel {
   } as LanguageModel;
 }
 
-function createTextOnlyMockModel(): LanguageModel {
+function createTextOnlyMockModel(
+  recordPromptRoles: (roles: string[]) => void,
+  shouldOverflow: () => boolean
+): LanguageModel {
   return {
     specificationVersion: "v3",
     provider: "test",
@@ -639,10 +643,33 @@ function createTextOnlyMockModel(): LanguageModel {
     doGenerate() {
       throw new Error("doGenerate not implemented");
     },
-    doStream() {
+    doStream(options: Record<string, unknown>) {
+      const prompt = Array.isArray(options.prompt) ? options.prompt : [];
+      recordPromptRoles(
+        prompt.flatMap((message) => {
+          if (
+            typeof message !== "object" ||
+            message === null ||
+            !("role" in message) ||
+            typeof message.role !== "string"
+          ) {
+            return [];
+          }
+          return [message.role];
+        })
+      );
+
       const stream = new ReadableStream({
         start(controller) {
           controller.enqueue({ type: "stream-start", warnings: [] });
+          if (shouldOverflow()) {
+            controller.enqueue({
+              type: "error",
+              error: new Error("context_length_exceeded")
+            });
+            controller.close();
+            return;
+          }
           controller.enqueue({ type: "text-start", id: "t1" });
           controller.enqueue({
             type: "text-delta",
@@ -683,6 +710,26 @@ export class ThinkClientToolsAgent extends Think {
   private _slowChunkCount = 4;
   private _responseLog: ChatResponseResult[] = [];
   private _lastTurnToolNames: string[] = [];
+  private _textOnlyPromptRoles: string[][] = [];
+  private _textOnlyOverflowAttemptsRemaining = 0;
+  private _compactionHistoryMessageIds: string[][] = [];
+
+  override classifyChatError = defaultContextOverflowClassifier;
+
+  override configureSession(session: Session): Session {
+    return session.onCompaction(async (messages) => {
+      this._compactionHistoryMessageIds.push(
+        messages.map((message) => message.id)
+      );
+      const first = messages[0];
+      if (!first) return null;
+      return {
+        summary: "Compacted selected regeneration branch",
+        fromMessageId: first.id,
+        toMessageId: first.id
+      };
+    });
+  }
 
   override beforeTurn(ctx: { tools: ToolSet }): void {
     this._lastTurnToolNames = Object.keys(ctx.tools);
@@ -726,7 +773,15 @@ export class ThinkClientToolsAgent extends Think {
         this._midStreamParallelGapsBeforeSlow,
         this._midStreamParallelGapsAfterSlow
       );
-    if (this._useTextOnly) return createTextOnlyMockModel();
+    if (this._useTextOnly)
+      return createTextOnlyMockModel(
+        (roles) => this._textOnlyPromptRoles.push(roles),
+        () => {
+          if (this._textOnlyOverflowAttemptsRemaining === 0) return false;
+          this._textOnlyOverflowAttemptsRemaining--;
+          return true;
+        }
+      );
     if (this._useServerApprovalTool) return createServerApprovalToolMockModel();
     return createClientToolMockModel();
   }
@@ -757,6 +812,23 @@ export class ThinkClientToolsAgent extends Think {
     this._useTextOnly = value;
   }
 
+  /** Model prompt roles recorded by the text-only boundary implementation. */
+  async getTextOnlyPromptRoles(): Promise<string[][]> {
+    return this._textOnlyPromptRoles;
+  }
+
+  /** Make the next text-only attempt overflow and enable compact-and-retry. */
+  async overflowNextTextOnlyAttempt(): Promise<void> {
+    this.contextOverflow = { reactive: true };
+    this._textOnlyOverflowAttemptsRemaining = 1;
+    this._compactionHistoryMessageIds = [];
+  }
+
+  /** Message IDs supplied to each compaction function invocation. */
+  async getCompactionHistoryMessageIds(): Promise<string[][]> {
+    return this._compactionHistoryMessageIds;
+  }
+
   async setServerApprovalToolMode(value: boolean): Promise<void> {
     this._useServerApprovalTool = value;
   }