fix(session): remap compaction tail_start_id when forking (#24898)

Co-authored-by: spark4862 <[email protected]>
Co-authored-by: Aiden Cline <[email protected]>

2 files changed, 75 insertions, 2 deletions

diff --git a/packages/opencode/src/session/session.ts b/packages/opencode/src/session/session.ts
index 45b8f0078..c376c8d1a 100644
--- a/packages/opencode/src/session/session.ts
+++ b/packages/opencode/src/session/session.ts
@@ -616,12 +616,16 @@ export const layer: Layer.Layer<Service, never, Bus.Service | Storage.Service> =
         })
 
         for (const part of msg.parts) {
-          yield* updatePart({
+          const p: MessageV2.Part = {
             ...part,
             id: PartID.ascending(),
             messageID: cloned.id,
             sessionID: session.id,
-          })
+          }
+          if (p.type === "compaction" && p.tail_start_id) {
+            p.tail_start_id = idMap.get(p.tail_start_id)
+          }
+          yield* updatePart(p)
         }
       }
       return session
diff --git a/packages/opencode/test/session/messages-pagination.test.ts b/packages/opencode/test/session/messages-pagination.test.ts
index 3fe64ce80..17370bbe6 100644
--- a/packages/opencode/test/session/messages-pagination.test.ts
+++ b/packages/opencode/test/session/messages-pagination.test.ts
@@ -29,6 +29,9 @@ const svc = {
   updatePart<T extends MessageV2.Part>(part: T) {
     return run(SessionNs.Service.use((svc) => svc.updatePart(part)))
   },
+  fork(input: { sessionID: SessionID; messageID?: MessageID }) {
+    return run(SessionNs.Service.use((svc) => svc.fork(input)))
+  },
 }
 
 async function fill(sessionID: SessionID, count: number, time = (i: number) => Date.now() + i) {
@@ -837,6 +840,72 @@ describe("MessageV2.filterCompacted", () => {
     })
   })
 
+  test("fork remaps compaction tail_start_id for filterCompacted", async () => {
+    await Instance.provide({
+      directory: root,
+      fn: async () => {
+        const session = await svc.create({})
+
+        const u1 = await addUser(session.id, "first")
+        const a1 = await addAssistant(session.id, u1, { finish: "end_turn" })
+        await svc.updatePart({
+          id: PartID.ascending(),
+          sessionID: session.id,
+          messageID: a1,
+          type: "text",
+          text: "first reply",
+        })
+
+        const u2 = await addUser(session.id, "second")
+        const a2 = await addAssistant(session.id, u2, { finish: "end_turn" })
+        await svc.updatePart({
+          id: PartID.ascending(),
+          sessionID: session.id,
+          messageID: a2,
+          type: "text",
+          text: "second reply",
+        })
+
+        const c1 = await addUser(session.id)
+        await addCompactionPart(session.id, c1, u2)
+        const s1 = await addAssistant(session.id, c1, { summary: true, finish: "end_turn" })
+        await svc.updatePart({
+          id: PartID.ascending(),
+          sessionID: session.id,
+          messageID: s1,
+          type: "text",
+          text: "summary",
+        })
+
+        const u3 = await addUser(session.id, "third")
+        const a3 = await addAssistant(session.id, u3, { finish: "end_turn" })
+        await svc.updatePart({
+          id: PartID.ascending(),
+          sessionID: session.id,
+          messageID: a3,
+          type: "text",
+          text: "third reply",
+        })
+
+        const parentFiltered = MessageV2.filterCompacted(MessageV2.stream(session.id))
+        expect(parentFiltered.map((item) => item.info.id)).toEqual([u2, a2, c1, s1, u3, a3])
+
+        const forked = await svc.fork({ sessionID: session.id })
+        const childFiltered = MessageV2.filterCompacted(MessageV2.stream(forked.id))
+        expect(childFiltered).toHaveLength(parentFiltered.length)
+
+        const tailPart = childFiltered.flatMap((m) => m.parts).find((p) => p.type === "compaction")
+        expect(tailPart?.type).toBe("compaction")
+        if (!tailPart || tailPart.type !== "compaction") throw new Error("Expected forked compaction part")
+        expect(tailPart.tail_start_id).toBeDefined()
+        expect(childFiltered.some((m) => m.info.id === tailPart.tail_start_id)).toBe(true)
+
+        await svc.remove(forked.id)
+        await svc.remove(session.id)
+      },
+    })
+  })
+
   test("retains an assistant tail when compaction starts inside a turn", async () => {
     await Instance.provide({
       directory: root,