fix(app): restore optimistic timeline state (#38693)

packages/app/src/pages/session/timeline/message-timeline.tsx

@@ -332,6 +332,7 @@ export function MessageTimeline(props: {
   const showHeader = createMemo(() => !!(titleValue() || parentID()))
   const projection = createTimelineProjection({
     messages: sessionMessages,
+    userMessages: () => props.userMessages,
     sessionMessages: projectedMessages,
     parts: getMsgParts,
     status: sessionStatus,

packages/app/src/pages/session/timeline/projection.ts

@@ -8,6 +8,7 @@ export { reuseTimelineRows } from "./row-reconciliation"
 
 export function createTimelineProjection(input: {
   messages: Accessor<Message[]>
+  userMessages: Accessor<UserMessage[]>
   sessionMessages: Accessor<SessionMessageInfo[]>
   parts: (messageID: string) => Part[]
   status: Accessor<SessionStatus>
@@ -36,6 +37,7 @@ export function createTimelineProjection(input: {
       input.showReasoningSummaries(),
       input.status().type,
       input.inlineComments(),
+      input.userMessages(),
     ),
   )
   const activeMessageID = createMemo(() => projection().activeMessageID)

packages/app/src/pages/session/timeline/rows-current.test.ts

@@ -46,6 +46,7 @@ describe("current session timeline rows", () => {
       true,
       "busy",
       true,
+      normalized.messages.filter((message) => message.role === "user"),
     )
 
     expect(result.activeMessageID).toBe("msg_3")
@@ -81,6 +82,7 @@ describe("current session timeline rows", () => {
       true,
       "idle",
       true,
+      normalized.messages.filter((message) => message.role === "user"),
     )
 
     expect(result.activeMessageID).toBe("msg_shell")
@@ -121,6 +123,7 @@ describe("current session timeline rows", () => {
       true,
       "idle",
       true,
+      normalized.messages.filter((message) => message.role === "user"),
     )
 
     expect(result.rows.map(TimelineRow.key)).toEqual([
@@ -131,4 +134,37 @@ describe("current session timeline rows", () => {
       "assistant-part:msg_user_2:msg_assistant_2:text:0",
     ])
   })
+
+  test("renders an optimistic user turn and thinking before the protocol message arrives", () => {
+    const source = [
+      { id: "msg_1", type: "user", text: "existing", time: { created: 1 } },
+    ] satisfies SessionMessageInfo[]
+    const normalized = normalizeSessionMessages("ses_1", source)
+    const optimistic = {
+      id: "msg_2",
+      sessionID: "ses_1",
+      role: "user" as const,
+      time: { created: 2 },
+      agent: "build",
+      model: { modelID: "model", providerID: "provider" },
+    }
+    const result = Timeline.constructSessionMessageRows(
+      source,
+      (messageID) =>
+        messageID === optimistic.id ? optimistic : normalized.messages.find((message) => message.id === messageID),
+      () => [],
+      true,
+      "busy",
+      true,
+      [...normalized.messages.filter((message) => message.role === "user"), optimistic],
+    )
+
+    expect(result.activeMessageID).toBe(optimistic.id)
+    expect(result.rows.map(TimelineRow.key)).toEqual([
+      "user-message:msg_1",
+      "turn-gap:msg_2",
+      "user-message:msg_2",
+      "thinking:msg_2",
+    ])
+  })
 })

packages/app/src/pages/session/timeline/rows.ts

@@ -39,6 +39,7 @@ export namespace Timeline {
     showReasoning: boolean,
     status: SessionStatus["type"],
     inlineComments: boolean,
+    projectedUserMessages: UserMessage[],
   ) {
     const turns: { user: UserMessage; assistants: AssistantMessage[] }[] = []
     const turnByUserID = new Map<string, (typeof turns)[number]>()
@@ -70,6 +71,14 @@ export namespace Timeline {
       turns.push(turn)
       turnByUserID.set(user.id, turn)
     })
+    const latestUserMessageID = turns.at(-1)?.user.id
+    projectedUserMessages.forEach((user) => {
+      if (turnByUserID.has(user.id)) return
+      if (latestUserMessageID && user.id < latestUserMessageID) return
+      const turn = { user, assistants: [] }
+      turns.push(turn)
+      turnByUserID.set(user.id, turn)
+    })
     const activeMessageID = turns.at(-1)?.user.id
     return {
       activeMessageID,