diff --git a/src/client/src/chatMessages.test.ts b/src/client/src/chatMessages.test.ts index 1850992..e6504f6 100644 --- a/src/client/src/chatMessages.test.ts +++ b/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" }] }, + ]); + }); }); diff --git a/src/client/src/chatMessages.ts b/src/client/src/chatMessages.ts index a052ada..9ebfea9 100644 --- a/src/client/src/chatMessages.ts +++ b/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")); diff --git a/src/client/src/chatTranscript.test.ts b/src/client/src/chatTranscript.test.ts new file mode 100644 index 0000000..535e1a2 --- /dev/null +++ b/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" } }, + }, + ]); + }); +}); diff --git a/src/client/src/chatTranscript.ts b/src/client/src/chatTranscript.ts index 6cdcdb2..a023c23 100644 --- a/src/client/src/chatTranscript.ts +++ b/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[] { diff --git a/src/client/src/controllers/sessionController.ts b/src/client/src/controllers/sessionController.ts index 3c90bae..e19ce29 100644 --- a/src/client/src/controllers/sessionController.ts +++ b/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); } diff --git a/src/client/src/sessionSocket.ts b/src/client/src/sessionSocket.ts index fcd5cc5..72e59a2 100644 --- a/src/client/src/sessionSocket.ts +++ b/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 { diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index 288ef7b..56448fd 100644 --- a/src/server/sessions/piSessionService.ts +++ b/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 }; diff --git a/src/shared/apiTypes.ts b/src/shared/apiTypes.ts index 1183810..4971da4 100644 --- a/src/shared/apiTypes.ts +++ b/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 }