fix(session-ui): scope context row keys (#44043)

Co-authored-by: Hona <10430890+Hona@users.noreply.github.com>

packages/session-ui/src/timeline/projection.test.ts

@@ -2,13 +2,17 @@ import { describe, expect, test } from "bun:test"
 import type { ModelRef, SessionMessageInfo } from "@opencode-ai/client/promise"
 import { createTimelineProjection, reuseTimelineRows, TimelineRow, type PartGroup } from "./projection"
 
-const context = (key: string, partIDs: string[], userMessageID = "user-1") =>
+const context = (
+  key: string,
+  partIDs: string[],
+  identity: { userMessageID?: string; messageID?: string } = {},
+) =>
   new TimelineRow.AssistantPart({
-    userMessageID,
+    userMessageID: identity.userMessageID ?? "user-1",
     group: {
       key,
       type: "context",
-      refs: partIDs.map((partID) => ({ messageID: "assistant-1", partID })),
+      refs: partIDs.map((partID) => ({ messageID: identity.messageID ?? "assistant-1", partID })),
     } satisfies PartGroup,
     previousAssistantPart: false,
   })
@@ -62,11 +66,18 @@ describe("reuseTimelineRows", () => {
     },
     {
       name: "does not reuse context identity across user messages",
-      previous: [context("context:a", ["a", "b"], "user-1")],
-      rows: [context("context:b", ["b"], "user-2")],
+      previous: [context("context:a", ["a", "b"], { userMessageID: "user-1" })],
+      rows: [context("context:b", ["b"], { userMessageID: "user-2" })],
       expected: ["assistant-part:user-2:context:b"],
       reused: [],
     },
+    {
+      name: "does not reuse context identity across assistant messages",
+      previous: [context("context:assistant-1:a", ["a"], { messageID: "assistant-1" })],
+      rows: [context("context:assistant-2:a", ["a"], { messageID: "assistant-2" })],
+      expected: ["assistant-part:user-1:context:assistant-2:a"],
+      reused: [],
+    },
     {
       name: "reuses an unaffected ordinary row",
       previous: [user()],

packages/session-ui/src/timeline/projection.ts

@@ -312,7 +312,7 @@ export function reuseTimelineRows(previous: TimelineRow.TimelineRow[] | undefine
   const contextByPart = new Map<string, PriorContext>()
   previous.forEach((row, index) => {
     if (row._tag !== "AssistantPart" || row.group.type !== "context") return
-    row.group.refs.forEach((ref) => contextByPart.set(`${row.userMessageID}:${ref.partID}`, { index, row }))
+    row.group.refs.forEach((ref) => contextByPart.set(contextPartKey(row.userMessageID, ref), { index, row }))
   })
   const reserved = new Map<string, number>()
   rows.forEach((row, index) => {
@@ -407,7 +407,7 @@ function stabilizeContextKey(
 ) {
   if (row._tag !== "AssistantPart" || row.group.type !== "context") return row
   const existing = row.group.refs.reduce<PriorContext | undefined>((result, ref) => {
-    const candidate = contextByPart.get(`${row.userMessageID}:${ref.partID}`)
+    const candidate = contextByPart.get(contextPartKey(row.userMessageID, ref))
     if (!candidate) return result
     const key = TimelineRow.key(candidate.row)
     if (claimed.has(key)) return result
@@ -426,6 +426,10 @@ function stabilizeContextKey(
   })
 }
 
+function contextPartKey(userMessageID: string, ref: PartRef) {
+  return `${userMessageID}:${ref.messageID}:${ref.partID}`
+}
+
 function renderable(content: Content, showReasoning: boolean) {
   if (content.type === "text") return !!content.text.trim()
   if (content.type === "reasoning") return showReasoning && !!content.text.trim()
@@ -440,7 +444,7 @@ function groupContent(items: { messageID: string; partID: string; content: Conte
   const flush = () => {
     const first = context[0]
     if (!first) return
-    groups.push({ type: "context", key: `context:${first.partID}`, refs: context })
+    groups.push({ type: "context", key: `context:${first.messageID}:${first.partID}`, refs: context })
     context = []
   }
 

packages/session-ui/src/timeline/rows-current.test.ts

@@ -262,7 +262,7 @@ describe("current session timeline rows", () => {
     expect(groups).toEqual([
       {
         type: "context",
-        key: "context:tool_read",
+        key: "context:msg_assistant:tool_read",
         refs: [
           { messageID: "msg_assistant", partID: "tool_read" },
           { messageID: "msg_assistant", partID: "tool_grep" },
@@ -276,6 +276,39 @@ describe("current session timeline rows", () => {
     ])
   })
 
+  test("keeps context row keys unique when tool IDs repeat across assistant messages", () => {
+    const tool = (name: string) => ({
+      type: "tool" as const,
+      id: "tool_0",
+      name,
+      state: { status: "running" as const, input: {}, metadata: {} },
+      time: { created: 2 },
+    })
+    const assistant = (id: string, name: string) => ({
+      id,
+      type: "assistant" as const,
+      agent: "build",
+      model: { id: "model", providerID: "provider" },
+      content: [tool(name)],
+      time: { created: 2 },
+    })
+    const source = [
+      { id: "msg_user", type: "user", text: "inspect", time: { created: 1 } },
+      assistant("msg_assistant_1", "read"),
+      assistant("msg_assistant_2", "execute"),
+      assistant("msg_assistant_3", "grep"),
+    ] satisfies SessionMessageInfo[]
+
+    const keys = Timeline.constructSessionMessageRows(source, false, { type: "idle" }).rows.map(TimelineRow.key)
+
+    expect(keys).toEqual([
+      "user-message:msg_user",
+      "assistant-part:msg_user:context:msg_assistant_1:tool_0",
+      "assistant-part:msg_user:part:msg_assistant_2:tool_0",
+      "assistant-part:msg_user:context:msg_assistant_3:tool_0",
+    ])
+  })
+
   test("places a divider after interrupted output unless the turn compacts", () => {
     const messages = [
       { id: "msg_user", type: "user", text: "continue", time: { created: 1 } },