fix(think): preserve regeneration branch context - #2038
Conversation
🦋 Changeset detectedLatest commit: 935bbb7 The changes in this PR will be included in the next version bump. This PR includes changesets to release 3 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
agents
@cloudflare/ai-chat
@cloudflare/codemode
hono-agents
@cloudflare/shell
@cloudflare/think
@cloudflare/voice
@cloudflare/worker-bundler
commit: |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
24f92d0 to
48e57e5
Compare
48e57e5 to
935bbb7
Compare
ben-reitz
added
the
ready-for-review-and-merge
| // 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; |
There was a problem hiding this comment.
🟡 Compacting a second conversation branch silently un-compacts the first one
Summaries stored for one conversation branch are matched only by their starting message (the filter added at packages/agents/src/experimental/memory/session/session.ts:780-785 selects an overlay to extend), so once a second branch is summarized from the same starting point the first branch's summary stops being applied at all.
Impact: A conversation that was already shortened silently goes back to its full length, so the next request can blow the model's context window again even though compaction "succeeded".
Overlay selection matches on start id only and ignores whether the range ends on this path
compact(leafId) now compacts a specific root-to-leaf branch and stores an overlay (fromMessageId, toMessageId). Because createCompactFunction always starts the summary at the first non-pinned message (packages/agents/src/experimental/memory/utils/compaction-helpers.ts:564), two sibling branches compacted independently produce two overlays with the SAME fromMessageId but toMessageIds on different branches.
Rendering in packages/agents/src/experimental/memory/session/providers/agent.ts:539-546 picks matching[matching.length - 1] — the newest overlay with that fromMessageId — without checking that its toMessageId is on the current path. When it is not (endIdx === -1), the branch renders with NO overlay at all, so the older, correct overlay for that branch is never applied.
The new filter in Session.compact then also fails to see an overlay id in historyIds, and falls back to the raw from/to clause, so a subsequent compaction on that branch keeps re-adding overlays that will keep losing the tie-break.
A fix belongs in applyCompactions: choose the widest matching compaction whose toMessageId is actually present on the path.
Prompt for agents
Session.compact(leafId) now supports compacting a specific branch, which makes it possible for two sibling branches to hold compaction overlays that share the same fromMessageId (createCompactFunction always starts at the first non-pinned message). AgentSessionProvider.applyCompactions selects overlays by fromMessageId and picks matching[matching.length - 1] (the newest) without verifying that its toMessageId exists on the path being rendered. When the newest overlay belongs to a sibling branch, endIdx is -1 and NO overlay is applied, so the branch that does have a valid overlay silently renders uncompacted. Consider filtering candidate compactions to those whose toMessageId is present in the current path before choosing the widest/newest one, and add a test covering two branches compacted from the same starting message.
Was this helpful? React with 👍 or 👎 to provide feedback.
| persistOrphanedStream: (streamId) => | ||
| this._persistOrphanedStream(streamId), | ||
| this._persistOrphanedStream(streamId, this._activeTurnHistory?.leafId), |
There was a problem hiding this comment.
🟡 Rebuilt partial answer after hibernation can still be filed under the previous answer instead of beside it
The branch chosen for a regenerated answer is read from memory (this._activeTurnHistory?.leafId at packages/think/src/think.ts:15829) at a point that is only ever reached after the agent has restarted and that memory is empty, so a rebuilt partial answer is filed at the end of the current conversation rather than next to the previous answer.
Impact: A regeneration interrupted by a restart can reappear as a continuation of the old answer instead of an alternative to it, unless chat recovery happens to rebuild it first.
Why the in-memory selection is always undefined on this path
ResumeHandshake.handleResumeAck only calls persistOrphanedStream when replayChunks returns a stream id, which happens exclusively in the !this._isLive branch (packages/agents/src/chat/resumable-stream.ts:524-543) — i.e. the stream was restored from SQLite after hibernation and has no live reader. In that isolate _activeTurnHistory was never assigned (it is only set in _prepareInferenceInvocation, packages/think/src/think.ts:5632), so the argument is always undefined and _persistOrphanedStream appends at the latest leaf (the previous answer).
With the default chatRecovery = true the wake-time fiber recovery usually persists the orphan first (with snapshot.historyLeafId), masking this; with chatRecovery disabled there is no fiber, and the ACK path is the only persister, so the regenerated partial is attached under the previous answer. Carrying the branch leaf durably (e.g. on the stream metadata row, alongside message_id) would make this path work for real.
Prompt for agents
In Think._resumeHandshake the orphan-persist seam passes this._activeTurnHistory?.leafId as the parent for the reconstructed message. That callback only runs for an orphaned stream (ResumableStream.replayChunks returns a stream id only when !_isLive, i.e. the stream was restored from SQLite after hibernation), at which point the in-memory _activeTurnHistory is always undefined, so the branch parent is never applied and the partial is appended under the current active leaf (the previous answer) for a regeneration turn. Consider persisting the branch leaf durably with the stream (e.g. an extra column on cf_ai_chat_stream_metadata next to message_id, written in _startResumableStream) and reading it back in _persistOrphanedStream, so both the wake-recovery path and the resume-handshake path reconstruct the partial as a sibling branch.
Was this helpful? React with 👍 or 👎 to provide feedback.
1 participant
Bug
Regeneration means “answer this user message again” while keeping the previous answer available. Think already stored the new answer as a sibling, but it generated that answer from the wrong history.
The result looked like a new branch in storage, but the model had been asked to continue the old answer instead of producing an independent alternative.
Fix
Think now uses one server-validated branch point for both model input and output storage. For regeneration, model history ends at the selected user message, and the generated answer is stored beneath that same message.
The selected branch is also preserved through hooks, retries, context compaction, Durable Object eviction recovery, and partial-response reconstruction. Re-delivered recovery work is idempotent, so it cannot create duplicate alternatives.
Summary
Fixes #2028
Testing
pnpm run checkpnpm --dir packages/agents run test:chat— 514 passedpnpm --dir packages/think run test:workers— 923 passed