fix(app): preserve tool disclosures when groups update (#45474)

packages/app/e2e/regression/session-timeline-collapse-state.spec.ts

@@ -84,6 +84,39 @@ const assistantMessage = {
 } satisfies SessionMessageInfo
 
 test.describe("regression: session timeline local row state", () => {
+  test("preserves a patch file choice as new calls join its Used group", async ({ page }) => {
+    const events: EventPayload[] = []
+    const part = { ...editPart, tool: "patch" }
+    await mockServer(page, events, [userMessage, { ...assistantMessage, content: [toolContent(part)] }])
+    await configurePage(page, false)
+    await page.goto(sessionHref())
+    await expectSessionTitle(page, title)
+
+    const group = page.locator('[data-component="collapsed-tool-group"]')
+    const summary = group.getByRole("button", { name: "Used Patch", exact: true })
+    await summary.click()
+    await group.locator(`[data-timeline-part-id="${editPartID}"]`).evaluate((element) => {
+      element.setAttribute("data-disclosure-probe", "existing")
+    })
+    const wrapper = group.locator('[data-disclosure-probe="existing"]')
+    const trigger = wrapper.locator('[data-scope="apply-patch"] button')
+    await expect(trigger).toHaveAttribute("aria-expanded", "false")
+    await trigger.click()
+    await expect(trigger).toHaveAttribute("aria-expanded", "true")
+    const original = await wrapper.elementHandle()
+
+    for (const count of [2, 3]) {
+      if (count === 3) await trigger.click()
+      const id = `prt_patch_${count}`
+      events.push(...toolEvents({ ...part, id, callID: id }))
+      await expect(group.locator('[data-component="tag"]')).toHaveText(String(count))
+      await expect(group).toHaveAttribute("data-timeline-part-ids", new RegExp(`${id}$`))
+      await expect(trigger).toHaveAttribute("aria-expanded", String(count === 2))
+      await expect(summary).toHaveAttribute("aria-expanded", "true")
+      expect(await original!.evaluate((node) => node.isConnected)).toBe(true)
+    }
+  })
+
   test("keeps a manually collapsed tool collapsed when later assistant content streams", async ({ page }) => {
     const events: EventPayload[] = []
     await mockServer(page, events)
@@ -208,19 +241,19 @@ test.describe("regression: session timeline local row state", () => {
   })
 })
 
-async function configurePage(page: Page) {
-  await page.addInitScript(() => {
+async function configurePage(page: Page, expanded = true) {
+  await page.addInitScript((expanded) => {
     localStorage.setItem(
       "settings.v3",
       JSON.stringify({
         general: {
-          editToolPartsExpanded: true,
-          shellToolPartsExpanded: true,
+          editToolPartsExpanded: expanded,
+          shellToolPartsExpanded: expanded,
           showReasoningSummaries: true,
         },
       }),
     )
-  })
+  }, expanded)
 }
 
 async function expectExpanded(locator: Locator, expected: boolean) {

packages/session-ui/component-tests/tool-disclosure.spec.ts

@@ -0,0 +1,48 @@
+import { expect, story } from "../../storybook/playwright/story"
+
+for (const open of [true, false]) {
+  story(
+    `preserves ${open ? "expanded" : "collapsed"} tool choices when calls join the group`,
+    async ({ mount }, info) => {
+      const root = await mount("current-session-file-changes--appending-tool-calls")
+      const group = root.locator('[data-component="collapsed-tool-group"]')
+      await group.getByRole("button", { name: "Used Shell, Patch", exact: true }).click()
+      const shell = group.locator('[data-timeline-part-id="tool_shell_existing"] [data-slot="collapsible-trigger"]')
+    await group.locator('[data-timeline-part-id="tool_patch_existing"]').evaluate((element) => {
+      element.setAttribute("data-disclosure-probe", "existing")
+    })
+    const patch = group.locator('[data-disclosure-probe="existing"]')
+      const first = patch.locator('[data-scope="apply-patch"] button').filter({ hasText: "a.ts" })
+      const second = patch.locator('[data-scope="apply-patch"] button').filter({ hasText: "b.ts" })
+      const diff = patch.locator('[data-type="update"]').filter({ hasText: "b.ts" }).locator('[data-component="file"]')
+      await shell.click()
+      await first.click()
+      await second.click()
+      if (!open) {
+        await shell.click()
+        await first.click()
+      }
+      await expect(shell).toHaveAttribute("aria-expanded", String(open))
+      await expect(first).toHaveAttribute("aria-expanded", String(open))
+      await expect(second).toHaveAttribute("aria-expanded", "true")
+      await expect(diff).toBeVisible()
+      const original = await patch.elementHandle()
+      for (const count of [3, 4]) {
+        await root.getByRole("button", { name: "Append tool call", exact: true }).click()
+        await expect(group.locator('[data-component="tag"]')).toHaveText(String(count))
+        await expect(diff).toBeVisible()
+        await root
+          .locator('[data-component="session-timeline"]')
+          .screenshot({ path: info.outputPath(`append-${count}.png`) })
+        await expect(shell).toHaveAttribute("aria-expanded", String(open))
+        await expect(first).toHaveAttribute("aria-expanded", String(open))
+        await expect(second).toHaveAttribute("aria-expanded", "true")
+        expect(await original!.evaluate((node) => node.isConnected)).toBe(true)
+        await expect(group.getByRole("button", { name: "Used Shell, Patch", exact: true })).toHaveAttribute(
+          "aria-expanded",
+          "true",
+        )
+      }
+    },
+  )
+}

packages/session-ui/src/timeline/file-changes.stories.tsx

@@ -199,6 +199,51 @@ const fileScenarios = {
   write: WrittenSource,
 }
 
+export const AppendingToolCalls = {
+  render: () => {
+    const [state, setState] = createStore({ calls: 0 })
+    const files = ["src/a.ts", "src/b.ts"].map((file) => ({
+      ...storyPatchFile(file),
+      patch: createTwoFilesPatch(file, file, "export const before = true\n", "export const after = true\n"),
+    }))
+    const document = createMemo(() =>
+      storyDocument([
+        storyTool("tool_shell_existing", "shell", "completed", { command: "printf checked" }, { output: "checked" }),
+        storyTool(
+          "tool_patch_existing",
+          "patch",
+          "completed",
+          { patchText: "Update two files" },
+          {
+            metadata: { files },
+          },
+        ),
+        ...Array.from({ length: state.calls }, (_, index) =>
+          storyTool(
+            `tool_patch_next_${index}`,
+            "patch",
+            "completed",
+            { patchText: "Update src/a.ts again" },
+            {
+              metadata: { files: [files[0]] },
+            },
+          ),
+        ),
+      ]),
+    )
+    return (
+      <section class="mx-auto flex w-full max-w-[860px] flex-col gap-4 p-6">
+        <button type="button" onClick={() => setState("calls", (count) => count + 1)}>
+          Append tool call
+        </button>
+        <CurrentSessionProviders document={document()}>
+          <SessionTimeline document={document()} />
+        </CurrentSessionProviders>
+      </section>
+    )
+  },
+}
+
 export const ChangingFiles = {
   args: { scenario: "streaming" },
   argTypes: { scenario: { control: "select", options: Object.keys(fileScenarios) } },

packages/session-ui/src/timeline/session-timeline-row.tsx

@@ -475,11 +475,13 @@ export function createSessionTimelineRowRenderer(input: {
         if (value._tag !== "AssistantPart") throw new Error("Expected an assistant-part timeline row")
         return value
       }
+      // Construct once per row key, not inside JSX that reruns when group refs change.
+      const content = renderAssistant(current, onSizeChange)
       return (
         <Frame row={current()}>
           <div data-slot="session-turn-message-container" class={`w-full ${padding()}`}>
             <div data-slot="session-turn-assistant-content" aria-hidden={workingTurn(current().userMessageID)}>
-              {renderAssistant(current, onSizeChange)}
+              {content}
             </div>
           </div>
         </Frame>