Align live assistant transcript with history

src/client/src/chatMessages.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { appendText, normalizeMessage, normalizeMessages, textMessage } from "./chatMessages";
+import { appendText, appendThinking, normalizeMessage, normalizeMessages, textMessage } from "./chatMessages";
 
 describe("chat message normalization", () => {
   it("normalizes simple text messages and drops empty content", () => {
@@ -51,10 +51,30 @@ describe("appendText", () => {
     ]);
   });
 
-  it("starts a new message when role or last part does not match", () => {
+  it("starts a new message when role does not match", () => {
     expect(appendText([textMessage("user", "hello")], "assistant", "hi")).toEqual([
       textMessage("user", "hello"),
       textMessage("assistant", "hi"),
     ]);
   });
+
+  it("adds a text part to the previous same-role non-text message", () => {
+    expect(appendText([{ role: "assistant", parts: [{ type: "thinking", text: "plan" }] }], "assistant", "answer")).toEqual([
+      { role: "assistant", parts: [{ type: "thinking", text: "plan" }, { type: "text", text: "answer" }] },
+    ]);
+  });
+});
+
+describe("appendThinking", () => {
+  it("appends thinking deltas to the previous assistant thinking part", () => {
+    expect(appendThinking([{ role: "assistant", parts: [{ type: "thinking", text: "pla" }] }], "n")).toEqual([
+      { role: "assistant", parts: [{ type: "thinking", text: "plan" }] },
+    ]);
+  });
+
+  it("adds a thinking part to the previous assistant message", () => {
+    expect(appendThinking([textMessage("assistant", "answer")], "plan")).toEqual([
+      { role: "assistant", parts: [{ type: "text", text: "answer" }, { type: "thinking", text: "plan" }] },
+    ]);
+  });
 });

src/client/src/chatMessages.ts

@@ -14,6 +14,7 @@ export function withMessageMeta(line: ChatLine, rawMessage: unknown): ChatLine {
 }
 
 export function appendText(messages: ChatLine[], role: ChatLine["role"], text: string): ChatLine[] {
+  if (text === "") return messages;
   const last = messages.at(-1);
   const lastPart = last?.parts.at(-1);
   if (last?.role === role && lastPart?.type === "text") {
@@ -22,9 +23,24 @@ export function appendText(messages: ChatLine[], role: ChatLine["role"], text: s
       { ...last, parts: [...last.parts.slice(0, -1), { ...lastPart, text: lastPart.text + text }] },
     ];
   }
+  if (last?.role === role) return [...messages.slice(0, -1), { ...last, parts: [...last.parts, { type: "text", text }] }];
   return [...messages, textMessage(role, text)];
 }
 
+export function appendThinking(messages: ChatLine[], text: string): ChatLine[] {
+  if (text === "") return messages;
+  const last = messages.at(-1);
+  const lastPart = last?.parts.at(-1);
+  if (last?.role === "assistant" && lastPart?.type === "thinking") {
+    return [
+      ...messages.slice(0, -1),
+      { ...last, parts: [...last.parts.slice(0, -1), { ...lastPart, text: lastPart.text + text }] },
+    ];
+  }
+  if (last?.role === "assistant") return [...messages.slice(0, -1), { ...last, parts: [...last.parts, { type: "thinking", text }] }];
+  return [...messages, { role: "assistant", parts: [{ type: "thinking", text }] }];
+}
+
 export function normalizeMessage(message: unknown): ChatLine[] {
   if (getString(message, "role") === "bashExecution") return [withMessageMeta(normalizeBashExecution(message), message)];
   const role = normalizeRole(getString(message, "role"));

src/client/src/chatTranscript.test.ts

@@ -0,0 +1,44 @@
+import { describe, expect, it } from "vitest";
+import { textMessage } from "./chatMessages";
+import { applyTranscriptEvent } from "./chatTranscript";
+import type { ChatLine } from "./components/shared";
+
+const finalAssistant = {
+  role: "assistant",
+  content: [
+    { type: "thinking", thinking: "plan" },
+    { type: "text", text: "answer" },
+  ],
+  timestamp: "2026-05-09T12:00:00.000Z",
+  provider: "test",
+  model: "model",
+};
+
+describe("applyTranscriptEvent", () => {
+  it("streams thinking and text into one assistant message", () => {
+    let messages: ChatLine[] = [];
+    messages = applyTranscriptEvent(messages, { type: "assistant.thinking.delta", text: "pla" }) ?? messages;
+    messages = applyTranscriptEvent(messages, { type: "assistant.thinking.delta", text: "n" }) ?? messages;
+    messages = applyTranscriptEvent(messages, { type: "assistant.delta", text: "answer" }) ?? messages;
+
+    expect(messages).toEqual([
+      { role: "assistant", parts: [{ type: "thinking", text: "plan" }, { type: "text", text: "answer" }] },
+    ]);
+  });
+
+  it("replaces the streamed assistant message with the finalized history shape", () => {
+    const streamed: ChatLine[] = [
+      textMessage("user", "question"),
+      { role: "assistant", parts: [{ type: "thinking", text: "partial" }, { type: "text", text: "partial answer" }] },
+    ];
+
+    expect(applyTranscriptEvent(streamed, { type: "message.end", message: finalAssistant })).toEqual([
+      textMessage("user", "question"),
+      {
+        role: "assistant",
+        parts: [{ type: "thinking", text: "plan" }, { type: "text", text: "answer" }],
+        meta: { timestamp: "2026-05-09T12:00:00.000Z", model: { provider: "test", id: "model" } },
+      },
+    ]);
+  });
+});

src/client/src/chatTranscript.ts

@@ -1,4 +1,4 @@
-import { appendText, normalizeMessage, textMessage, withMessageMeta } from "./chatMessages";
+import { appendText, appendThinking, normalizeMessage, textMessage } from "./chatMessages";
 import type { ChatLine } from "./components/shared";
 import { appendShellChunk, finalizeShellMessage, shellStartMessage } from "./shellMessages";
 import type { SessionUiEvent } from "./sessionSocket";
@@ -6,6 +6,7 @@ import type { SessionUiEvent } from "./sessionSocket";
 export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent): ChatLine[] | undefined {
   if (event.type === "message.append") return appendNormalized(messages, event.message);
   if (event.type === "assistant.delta") return appendText(messages, "assistant", event.text);
+  if (event.type === "assistant.thinking.delta") return appendThinking(messages, event.text);
   if (event.type === "tool.start") return appendNormalized(messages, { role: "assistant", content: [{ type: "toolCall", name: event.toolName, arguments: event.args }] });
   if (event.type === "tool.end") return appendNormalized(messages, { role: "toolResult", toolName: event.toolName, content: event.content ?? [{ type: "text", text: event.text }], isError: event.isError });
   if (event.type === "shell.start") return [...messages, shellStartMessage(event.command, event.excludeFromContext)];
@@ -13,23 +14,16 @@ export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent
   if (event.type === "shell.end") return finalizeShellMessage(messages, event);
   if (event.type === "command.output") return [...messages, textMessage(event.level === "error" ? "system" : "tool", event.message)];
   if (event.type === "session.error") return [...messages, textMessage("system", event.message)];
-  if (event.type === "message.end") return event.message === undefined ? undefined : applyMessageEndMeta(messages, event.message);
+  if (event.type === "message.end") return event.message === undefined ? undefined : applyFinalMessage(messages, event.message);
   return undefined;
 }
 
-function applyMessageEndMeta(messages: ChatLine[], rawMessage: unknown): ChatLine[] | undefined {
+function applyFinalMessage(messages: ChatLine[], rawMessage: unknown): ChatLine[] | undefined {
   const ended = normalizeMessage(rawMessage)[0];
-  if (ended?.meta === undefined) return undefined;
-  const index = findLastMatchingRole(messages, ended.role);
-  if (index < 0) return undefined;
-  return messages.map((message, i) => i === index ? withMessageMeta(message, rawMessage) : message);
-}
-
-function findLastMatchingRole(messages: ChatLine[], role: ChatLine["role"]): number {
-  for (let i = messages.length - 1; i >= 0; i -= 1) {
-    if (messages[i]?.role === role) return i;
-  }
-  return -1;
+  if (ended === undefined) return undefined;
+  const last = messages.at(-1);
+  if (last?.role !== ended.role) return [...messages, ended];
+  return [...messages.slice(0, -1), ended];
 }
 
 function appendNormalized(messages: ChatLine[], rawMessage: unknown): ChatLine[] {

src/client/src/controllers/sessionController.ts

@@ -304,6 +304,6 @@ export class SessionController {
 }
 
 function isTranscriptEvent(event: SessionUiEvent): boolean {
-  return ["message.append", "assistant.delta", "tool.start", "tool.end", "shell.start", "shell.chunk", "shell.end", "command.output", "session.error"].includes(event.type);
+  return ["message.append", "assistant.delta", "assistant.thinking.delta", "tool.start", "tool.end", "shell.start", "shell.chunk", "shell.end", "command.output", "session.error"].includes(event.type);
 }
 

src/client/src/sessionSocket.ts

@@ -114,7 +114,7 @@ export class GlobalSessionSocket {
 
 function isSessionUiEvent(event: unknown): event is SessionUiEvent {
   const type = eventType(event);
-  return ["message.append", "assistant.delta", "tool.start", "tool.end", "shell.start", "shell.chunk", "shell.end", "agent.start", "agent.end", "message.end", "status.update", "activity.update", "command.output", "session.error", "session.name", "pi.event"].includes(type);
+  return ["message.append", "assistant.delta", "assistant.thinking.delta", "tool.start", "tool.end", "shell.start", "shell.chunk", "shell.end", "agent.start", "agent.end", "message.end", "status.update", "activity.update", "command.output", "session.error", "session.name", "pi.event"].includes(type);
 }
 
 function isGlobalSessionEvent(event: unknown): event is GlobalSessionEvent {

src/server/sessions/piSessionService.ts

@@ -440,6 +440,9 @@ function toClientEvent(event: unknown): SessionUiEvent {
   if (eventType === "message_update" && getString(assistantMessageEvent, "type") === "text_delta") {
     return { type: "assistant.delta", text: getString(assistantMessageEvent, "delta") ?? "" };
   }
+  if (eventType === "message_update" && getString(assistantMessageEvent, "type") === "thinking_delta") {
+    return { type: "assistant.thinking.delta", text: getString(assistantMessageEvent, "delta") ?? "" };
+  }
   if (eventType === "tool_execution_start") {
     const args = getProperty(event, "args");
     return { type: "tool.start", toolName: getString(event, "toolName") ?? "", toolCallId: getString(event, "toolCallId") ?? "", summary: summarizeToolArgs(args), args };

src/shared/apiTypes.ts

@@ -143,6 +143,7 @@ export type CommandResult =
 export type SessionUiEvent =
   | { type: "message.append"; message: unknown }
   | { type: "assistant.delta"; text: string }
+  | { type: "assistant.thinking.delta"; text: string }
   | { type: "tool.start"; toolName: string; toolCallId: string; summary: string; args?: unknown }
   | { type: "tool.end"; toolName: string; toolCallId: string; text: string; isError: boolean; content?: unknown }
   | { type: "shell.start"; command: string; excludeFromContext?: boolean }