From 95deb9b68a3a419ec16225c157ed8a495fbcab96 Mon Sep 17 00:00:00 2001
From: Kit Langton <kit.langton@gmail.com>
Date: Fri, 17 Jul 2026 23:08:05 -0400
Subject: [PATCH] fix(tui): preserve part collection identity across sync

---
 packages/tui/src/context/data.tsx          |  2 +-
 packages/tui/src/routes/session/content.ts | 12 +++
 packages/tui/test/cli/tui/data.test.tsx    | 96 ++++++++++++++++++++++
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/packages/tui/src/context/data.tsx b/packages/tui/src/context/data.tsx
index 3dc776beee..1fd768dfbb 100644
--- a/packages/tui/src/context/data.tsx
+++ b/packages/tui/src/context/data.tsx
@@ -955,7 +955,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
                 await client.api.message.list({ sessionID, limit: 200, order: "desc" })
               ).data.toReversed()
               messageIndex.set(sessionID, new Map(messages.map((message, index) => [message.id, index])))
-              content.drop(sessionID)
+              content.prune(sessionID, new Set(messages.map((message) => message.id)))
               for (const item of messages) {
                 if (item.type === "assistant") content.seed(sessionID, item.id, item.content)
               }
diff --git a/packages/tui/src/routes/session/content.ts b/packages/tui/src/routes/session/content.ts
index 4e4e737da9..8692879503 100644
--- a/packages/tui/src/routes/session/content.ts
+++ b/packages/tui/src/routes/session/content.ts
@@ -86,6 +86,18 @@ export namespace SessionContent {
           if (key.startsWith(prefix)) collections.delete(key)
         }
       },
+      /**
+       * Drop collections for messages absent from a refetched snapshot.
+       * Present messages must be reseeded in place, never dropped: mounted
+       * views hold their collection reference for their whole lifetime, so
+       * replacing the object would orphan their subscriptions.
+       */
+      prune(sessionID: string, keep: ReadonlySet<string>) {
+        const prefix = `${sessionID}\u0000`
+        for (const key of [...collections.keys()]) {
+          if (key.startsWith(prefix) && !keep.has(key.slice(prefix.length))) collections.delete(key)
+        }
+      },
     }
   }
 }
diff --git a/packages/tui/test/cli/tui/data.test.tsx b/packages/tui/test/cli/tui/data.test.tsx
index a14a842846..48470b0a03 100644
--- a/packages/tui/test/cli/tui/data.test.tsx
+++ b/packages/tui/test/cli/tui/data.test.tsx
@@ -973,6 +973,102 @@ test("publishes streaming deltas to one part slot without touching siblings", as
   }
 })
 
+test("keeps part slot references alive across message sync", async () => {
+  const events = createEventStream()
+  const sessionID = "session-slot-identity"
+  // The refetch snapshot is identical to the streamed state, so any failure
+  // here is collection identity loss, not snapshot staleness.
+  let snapshot: unknown[] = []
+  const calls = createFetch((url) => {
+    if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: snapshot, cursor: {} })
+  }, events)
+  let data!: ReturnType<typeof useData>
+  let client!: ReturnType<typeof useClient>
+
+  function Probe() {
+    data = useData()
+    client = useClient()
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <ClientProvider api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </ClientProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    await wait(() => client.connection.status() === "connected")
+    emitEvent(events, {
+      id: "evt_identity_step",
+      created: 1,
+      type: "session.step.started",
+      durable: durable(sessionID),
+      data: {
+        sessionID,
+        assistantMessageID: "message-assistant",
+        agent: "build",
+        model: { id: "model", providerID: "provider" },
+      },
+    })
+    emitEvent(events, {
+      id: "evt_identity_text",
+      created: 2,
+      type: "session.text.started",
+      durable: durable(sessionID, 1),
+      data: { sessionID, assistantMessageID: "message-assistant", ordinal: 0 },
+    })
+    emitEvent(events, {
+      id: "evt_identity_ended",
+      created: 3,
+      type: "session.text.ended",
+      durable: durable(sessionID, 2),
+      data: { sessionID, assistantMessageID: "message-assistant", ordinal: 0, text: "hello world" },
+    })
+    // A mounted view captures the collection once, at useSlot setup.
+    const mounted = data.session.message.parts(sessionID, "message-assistant")
+    await wait(() => {
+      const part = mounted.get("text:0")?.()
+      return part?.type === "text" && part.text === "hello world"
+    })
+
+    snapshot = [
+      {
+        id: "message-assistant",
+        type: "assistant",
+        agent: "build",
+        model: { id: "model", providerID: "provider" },
+        content: [{ type: "text", text: "hello world" }],
+        time: { created: 1 },
+      },
+    ]
+    await data.session.message.sync(sessionID)
+
+    // The registry must hand back the same collection the view is holding.
+    expect(data.session.message.parts(sessionID, "message-assistant")).toBe(mounted)
+
+    emitEvent(events, {
+      id: "evt_identity_delta",
+      created: 4,
+      type: "session.text.delta",
+      data: { sessionID, assistantMessageID: "message-assistant", ordinal: 0, delta: "!!" },
+    })
+    // The delta must reach the reference a mounted view is subscribed to.
+    await wait(() => {
+      const part = mounted.get("text:0")?.()
+      return part?.type === "text" && part.text === "hello world!!"
+    })
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("does not publish timeline rows for duplicate streaming deltas", async () => {
   const events = createEventStream()
   const sessionID = "session-stream-metrics"