From b921298af7c31b3ce4916037f02348841514c2ae Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sun, 28 Dec 2025 14:12:08 +0100 Subject: [PATCH] Use exhaustive switch on message.role throughout coding-agent - addMessageToChat: exhaustive switch for all AgentMessage roles - renderSessionContext: delegates to addMessageToChat, special handling for assistant tool calls and tool results - export-html formatMessage: exhaustive switch for all AgentMessage roles - Removed isHookMessage, isBashExecutionMessage type guards in favor of role checks - Fixed imports and removed unused getLatestCompactionEntry --- .../coding-agent/src/core/agent-session.ts | 15 +- packages/coding-agent/src/core/compaction.ts | 130 +++++++----- packages/coding-agent/src/core/export-html.ts | 193 ++++++++++-------- packages/coding-agent/src/core/messages.ts | 145 +++++++++---- .../coding-agent/src/core/session-manager.ts | 91 ++++----- packages/coding-agent/src/index.ts | 3 - ...ction.ts => compaction-summary-message.ts} | 19 +- .../interactive/components/user-message.ts | 8 +- .../src/modes/interactive/interactive-mode.ts | 187 ++++++++--------- packages/coding-agent/test/compaction.test.ts | 16 -- .../session-manager/build-context.test.ts | 11 +- 11 files changed, 442 insertions(+), 376 deletions(-) rename packages/coding-agent/src/modes/interactive/components/{compaction.ts => compaction-summary-message.ts} (63%) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 03806c51..9502a2a3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -34,7 +34,7 @@ import type { TurnEndEvent, TurnStartEvent, } from "./hooks/index.js"; -import { type BashExecutionMessage, type HookMessage, isHookMessage } from "./messages.js"; +import type { BashExecutionMessage, HookMessage } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; import type { CompactionEntry, SessionManager } from "./session-manager.js"; import type { SettingsManager, SkillsSettings } from "./settings-manager.js"; @@ -218,8 +218,8 @@ export class AgentSession { // Handle session persistence if (event.type === "message_end") { - // Check if this is a hook message (has _hookData marker) - if (isHookMessage(event.message)) { + // Check if this is a hook message + if (event.message.role === "hookMessage") { // Persist as CustomMessageEntry this.sessionManager.appendCustomMessageEntry( event.message.customType, @@ -227,10 +227,15 @@ export class AgentSession { event.message.display, event.message.details, ); - } else { - // Regular message - persist as SessionMessageEntry + } else if ( + event.message.role === "user" || + event.message.role === "assistant" || + event.message.role === "toolResult" + ) { + // Regular LLM message - persist as SessionMessageEntry this.sessionManager.appendMessage(event.message); } + // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere // Track assistant message for auto-compaction (checked on agent_end) if (event.message.role === "assistant") { diff --git a/packages/coding-agent/src/core/compaction.ts b/packages/coding-agent/src/core/compaction.ts index 38e2dece..ac750a9e 100644 --- a/packages/coding-agent/src/core/compaction.ts +++ b/packages/coding-agent/src/core/compaction.ts @@ -8,8 +8,8 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai"; import { complete } from "@mariozechner/pi-ai"; -import { convertToLlm } from "./messages.js"; -import { type CompactionEntry, createSummaryMessage, type SessionEntry } from "./session-manager.js"; +import { convertToLlm, createBranchSummaryMessage, createHookMessage } from "./messages.js"; +import type { CompactionEntry, SessionEntry } from "./session-manager.js"; /** * Extract AgentMessage from an entry if it produces one. @@ -20,14 +20,10 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | null { return entry.message; } if (entry.type === "custom_message") { - return { - role: "user", - content: entry.content, - timestamp: new Date(entry.timestamp).getTime(), - }; + return createHookMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); } if (entry.type === "branch_summary") { - return createSummaryMessage(entry.summary, entry.timestamp); + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); } return null; } @@ -116,59 +112,65 @@ export function shouldCompact(contextTokens: number, contextWindow: number, sett export function estimateTokens(message: AgentMessage): number { let chars = 0; - // Handle bashExecution messages - if (message.role === "bashExecution") { - const bash = message as unknown as { command: string; output: string }; - chars = bash.command.length + bash.output.length; - return Math.ceil(chars / 4); - } - - // Handle user messages - if (message.role === "user") { - const content = (message as { content: string | Array<{ type: string; text?: string }> }).content; - if (typeof content === "string") { - chars = content.length; - } else if (Array.isArray(content)) { - for (const block of content) { - if (block.type === "text" && block.text) { - chars += block.text.length; + switch (message.role) { + case "user": { + const content = (message as { content: string | Array<{ type: string; text?: string }> }).content; + if (typeof content === "string") { + chars = content.length; + } else if (Array.isArray(content)) { + for (const block of content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } } } + return Math.ceil(chars / 4); } - return Math.ceil(chars / 4); - } - - // Handle assistant messages - if (message.role === "assistant") { - const assistant = message as AssistantMessage; - for (const block of assistant.content) { - if (block.type === "text") { - chars += block.text.length; - } else if (block.type === "thinking") { - chars += block.thinking.length; - } else if (block.type === "toolCall") { - chars += block.name.length + JSON.stringify(block.arguments).length; + case "assistant": { + const assistant = message as AssistantMessage; + for (const block of assistant.content) { + if (block.type === "text") { + chars += block.text.length; + } else if (block.type === "thinking") { + chars += block.thinking.length; + } else if (block.type === "toolCall") { + chars += block.name.length + JSON.stringify(block.arguments).length; + } } + return Math.ceil(chars / 4); } - return Math.ceil(chars / 4); - } - - // Handle tool results - if (message.role === "toolResult") { - const toolResult = message as { content: Array<{ type: string; text?: string }> }; - for (const block of toolResult.content) { - if (block.type === "text" && block.text) { - chars += block.text.length; + case "hookMessage": + case "toolResult": { + if (typeof message.content === "string") { + chars = message.content.length; + } else { + for (const block of message.content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } + if (block.type === "image") { + chars += 4800; // Estimate images as 4000 chars, or 1200 tokens + } + } } + return Math.ceil(chars / 4); + } + case "bashExecution": { + chars = message.command.length + message.output.length; + return Math.ceil(chars / 4); + } + case "branchSummary": + case "compactionSummary": { + chars = message.summary.length; + return Math.ceil(chars / 4); } - return Math.ceil(chars / 4); } return 0; } /** - * Find valid cut points: indices of user, assistant, or bashExecution messages. + * Find valid cut points: indices of user, assistant, custom, or bashExecution messages. * Never cut at tool results (they must follow their tool call). * When we cut at an assistant message with tool calls, its tool results follow it * and will be kept. @@ -178,16 +180,34 @@ function findValidCutPoints(entries: SessionEntry[], startIndex: number, endInde const cutPoints: number[] = []; for (let i = startIndex; i < endIndex; i++) { const entry = entries[i]; + switch (entry.type) { + case "message": { + const role = entry.message.role; + switch (role) { + case "bashExecution": + case "hookMessage": + case "branchSummary": + case "compactionSummary": + case "user": + case "assistant": + cutPoints.push(i); + break; + case "toolResult": + break; + } + break; + } + case "thinking_level_change": + case "model_change": + case "compaction": + case "branch_summary": + case "custom": + case "custom_message": + case "label": + } // branch_summary and custom_message are user-role messages, valid cut points if (entry.type === "branch_summary" || entry.type === "custom_message") { cutPoints.push(i); - } else if (entry.type === "message") { - const role = entry.message.role; - // user, assistant, and bashExecution are valid cut points - // toolResult must stay with its preceding tool call - if (role === "user" || role === "assistant" || role === "bashExecution") { - cutPoints.push(i); - } } } return cutPoints; diff --git a/packages/coding-agent/src/core/export-html.ts b/packages/coding-agent/src/core/export-html.ts index db53d743..f553e140 100644 --- a/packages/coding-agent/src/core/export-html.ts +++ b/packages/coding-agent/src/core/export-html.ts @@ -1,4 +1,4 @@ -import type { AgentState } from "@mariozechner/pi-agent-core"; +import type { AgentMessage, AgentState } from "@mariozechner/pi-agent-core"; import type { AssistantMessage, ImageContent, Message, ToolResultMessage, UserMessage } from "@mariozechner/pi-ai"; import { existsSync, readFileSync, writeFileSync } from "fs"; import hljs from "highlight.js"; @@ -7,7 +7,6 @@ import { homedir } from "os"; import * as path from "path"; import { basename } from "path"; import { APP_NAME, getCustomThemesDir, getThemesDir, VERSION } from "../config.js"; -import { type BashExecutionMessage, isBashExecutionMessage } from "./messages.js"; import type { SessionManager } from "./session-manager.js"; // ============================================================================ @@ -821,110 +820,136 @@ function formatToolExecution( return { html, bgColor }; } -function formatMessage(message: Message, toolResultsMap: Map, colors: ThemeColors): string { +function formatMessage( + message: AgentMessage, + toolResultsMap: Map, + colors: ThemeColors, +): string { let html = ""; const timestamp = (message as { timestamp?: number }).timestamp; const timestampHtml = timestamp ? `
${formatTimestamp(timestamp)}
` : ""; - // Handle bash execution messages (user-executed via ! command) - if (isBashExecutionMessage(message)) { - const bashMsg = message as unknown as BashExecutionMessage; - const isError = bashMsg.cancelled || (bashMsg.exitCode !== 0 && bashMsg.exitCode !== null); + switch (message.role) { + case "bashExecution": { + const isError = message.cancelled || (message.exitCode !== 0 && message.exitCode !== null); - html += `
`; - html += timestampHtml; - html += `
$ ${escapeHtml(bashMsg.command)}
`; + html += `
`; + html += timestampHtml; + html += `
$ ${escapeHtml(message.command)}
`; - if (bashMsg.output) { - const lines = bashMsg.output.split("\n"); - html += formatExpandableOutput(lines, 10); + if (message.output) { + const lines = message.output.split("\n"); + html += formatExpandableOutput(lines, 10); + } + + if (message.cancelled) { + html += `
(cancelled)
`; + } else if (message.exitCode !== 0 && message.exitCode !== null) { + html += `
(exit ${message.exitCode})
`; + } + + if (message.truncated && message.fullOutputPath) { + html += `
Output truncated. Full output: ${escapeHtml(message.fullOutputPath)}
`; + } + + html += `
`; + break; } + case "user": { + const userMsg = message as UserMessage; + let textContent = ""; + const images: ImageContent[] = []; - if (bashMsg.cancelled) { - html += `
(cancelled)
`; - } else if (bashMsg.exitCode !== 0 && bashMsg.exitCode !== null) { - html += `
(exit ${bashMsg.exitCode})
`; - } - - if (bashMsg.truncated && bashMsg.fullOutputPath) { - html += `
Output truncated. Full output: ${escapeHtml(bashMsg.fullOutputPath)}
`; - } - - html += `
`; - return html; - } - - if (message.role === "user") { - const userMsg = message as UserMessage; - let textContent = ""; - const images: ImageContent[] = []; - - if (typeof userMsg.content === "string") { - textContent = userMsg.content; - } else { - for (const block of userMsg.content) { - if (block.type === "text") { - textContent += block.text; - } else if (block.type === "image") { - images.push(block as ImageContent); + if (typeof userMsg.content === "string") { + textContent = userMsg.content; + } else { + for (const block of userMsg.content) { + if (block.type === "text") { + textContent += block.text; + } else if (block.type === "image") { + images.push(block as ImageContent); + } } } - } - html += `
${timestampHtml}`; + html += `
${timestampHtml}`; - // Render images first - if (images.length > 0) { - html += `
`; - for (const img of images) { - html += `User uploaded image`; + // Render images first + if (images.length > 0) { + html += `
`; + for (const img of images) { + html += `User uploaded image`; + } + html += `
`; } + + // Render text as markdown (server-side) + if (textContent.trim()) { + html += `
${renderMarkdown(textContent)}
`; + } + html += `
`; + break; } + case "assistant": { + html += timestampHtml ? `
${timestampHtml}` : ""; - // Render text as markdown (server-side) - if (textContent.trim()) { - html += `
${renderMarkdown(textContent)}
`; - } - - html += `
`; - } else if (message.role === "assistant") { - const assistantMsg = message as AssistantMessage; - html += timestampHtml ? `
${timestampHtml}` : ""; - - for (const content of assistantMsg.content) { - if (content.type === "text" && content.text.trim()) { - // Render markdown server-side - html += `
${renderMarkdown(content.text)}
`; - } else if (content.type === "thinking" && content.thinking.trim()) { - html += `
${escapeHtml(content.thinking.trim()).replace(/\n/g, "
")}
`; + for (const content of message.content) { + if (content.type === "text" && content.text.trim()) { + // Render markdown server-side + html += `
${renderMarkdown(content.text)}
`; + } else if (content.type === "thinking" && content.thinking.trim()) { + html += `
${escapeHtml(content.thinking.trim()).replace(/\n/g, "
")}
`; + } } - } - for (const content of assistantMsg.content) { - if (content.type === "toolCall") { - const toolResult = toolResultsMap.get(content.id); - const { html: toolHtml, bgColor } = formatToolExecution( - content.name, - content.arguments as Record, - toolResult, - colors, - ); - html += `
${toolHtml}
`; + for (const content of message.content) { + if (content.type === "toolCall") { + const toolResult = toolResultsMap.get(content.id); + const { html: toolHtml, bgColor } = formatToolExecution( + content.name, + content.arguments as Record, + toolResult, + colors, + ); + html += `
${toolHtml}
`; + } } - } - const hasToolCalls = assistantMsg.content.some((c) => c.type === "toolCall"); - if (!hasToolCalls) { - if (assistantMsg.stopReason === "aborted") { - html += '
Aborted
'; - } else if (assistantMsg.stopReason === "error") { - html += `
Error: ${escapeHtml(assistantMsg.errorMessage || "Unknown error")}
`; + const hasToolCalls = message.content.some((c) => c.type === "toolCall"); + if (!hasToolCalls) { + if (message.stopReason === "aborted") { + html += '
Aborted
'; + } else if (message.stopReason === "error") { + html += `
Error: ${escapeHtml(message.errorMessage || "Unknown error")}
`; + } } - } - if (timestampHtml) { - html += "
"; + if (timestampHtml) { + html += "
"; + } + break; + } + case "toolResult": + // Tool results are rendered inline with tool calls + break; + case "hookMessage": + // Hook messages with display:true shown as info boxes + if (message.display) { + const content = typeof message.content === "string" ? message.content : JSON.stringify(message.content); + html += `
${timestampHtml}
[${escapeHtml(message.customType)}]
${renderMarkdown(content)}
`; + } + break; + case "compactionSummary": + // Rendered separately via formatCompaction + break; + case "branchSummary": + // Rendered as compaction-like summary + html += `
Branch Summary
${escapeHtml(message.summary).replace(/\n/g, "
")}
`; + break; + default: { + // Exhaustive check + const _exhaustive: never = message; } } diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index 74451907..f9d33bbf 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -8,6 +8,21 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { ImageContent, Message, TextContent } from "@mariozechner/pi-ai"; +export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary: + + +`; + +export const COMPACTION_SUMMARY_SUFFIX = ` +`; + +export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from: + + +`; + +export const BRANCH_SUMMARY_SUFFIX = ``; + /** * Message type for bash executions via the ! command. */ @@ -35,28 +50,30 @@ export interface HookMessage { timestamp: number; } +export interface BranchSummaryMessage { + role: "branchSummary"; + summary: string; + fromId: string; + timestamp: number; +} + +export interface CompactionSummaryMessage { + role: "compactionSummary"; + summary: string; + tokensBefore: number; + timestamp: number; +} + // Extend CustomAgentMessages via declaration merging declare module "@mariozechner/pi-agent-core" { interface CustomAgentMessages { bashExecution: BashExecutionMessage; hookMessage: HookMessage; + branchSummary: BranchSummaryMessage; + compactionSummary: CompactionSummaryMessage; } } -/** - * Type guard for BashExecutionMessage. - */ -export function isBashExecutionMessage(msg: AgentMessage | Message): msg is BashExecutionMessage { - return msg.role === "bashExecution"; -} - -/** - * Type guard for HookMessage. - */ -export function isHookMessage(msg: AgentMessage | Message): msg is HookMessage { - return msg.role === "hookMessage"; -} - /** * Convert a BashExecutionMessage to user message text for LLM context. */ @@ -78,6 +95,46 @@ export function bashExecutionToText(msg: BashExecutionMessage): string { return text; } +export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage { + return { + role: "branchSummary", + summary, + fromId, + timestamp: new Date(timestamp).getTime(), + }; +} + +export function createCompactionSummaryMessage( + summary: string, + tokensBefore: number, + timestamp: string, +): CompactionSummaryMessage { + return { + role: "compactionSummary", + summary: summary, + tokensBefore, + timestamp: new Date(timestamp).getTime(), + }; +} + +/** Convert CustomMessageEntry to AgentMessage format */ +export function createHookMessage( + customType: string, + content: string | (TextContent | ImageContent)[], + display: boolean, + details: unknown | undefined, + timestamp: string, +): HookMessage { + return { + role: "hookMessage", + customType, + content, + display, + details, + timestamp: new Date(timestamp).getTime(), + }; +} + /** * Transform AgentMessages (including custom types) to LLM-compatible Messages. * @@ -89,30 +146,44 @@ export function bashExecutionToText(msg: BashExecutionMessage): string { export function convertToLlm(messages: AgentMessage[]): Message[] { return messages .map((m): Message | null => { - if (isBashExecutionMessage(m)) { - // Convert bash execution to user message - return { - role: "user", - content: [{ type: "text", text: bashExecutionToText(m) }], - timestamp: m.timestamp, - }; + switch (m.role) { + case "bashExecution": + return { + role: "user", + content: [{ type: "text", text: bashExecutionToText(m) }], + timestamp: m.timestamp, + }; + case "hookMessage": { + const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; + return { + role: "user", + content, + timestamp: m.timestamp, + }; + } + case "branchSummary": + return { + role: "user", + content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }], + timestamp: m.timestamp, + }; + case "compactionSummary": + return { + role: "user", + content: [ + { type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX }, + ], + timestamp: m.timestamp, + }; + case "user": + case "assistant": + case "toolResult": + return m; + default: + // biome-ignore lint/correctness/noSwitchDeclarations: fine + const _exhaustiveCheck: never = m; + return null; } - if (isHookMessage(m)) { - // Convert hook message to user message for LLM - // Normalize string content to array format - const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; - return { - role: "user", - content, - timestamp: m.timestamp, - }; - } - // Pass through standard LLM roles - if (m.role === "user" || m.role === "assistant" || m.role === "toolResult") { - return m as Message; - } - // Filter out unknown message types - return null; }) .filter((m) => m !== null); } diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 68a8e338..7035c1e0 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; -import type { ImageContent, TextContent } from "@mariozechner/pi-ai"; +import type { ImageContent, Message, TextContent } from "@mariozechner/pi-ai"; import { randomUUID } from "crypto"; import { appendFileSync, @@ -15,6 +15,13 @@ import { } from "fs"; import { join, resolve } from "path"; import { getAgentDir as getDefaultAgentDir } from "../config.js"; +import { + type BashExecutionMessage, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createHookMessage, + type HookMessage, +} from "./messages.js"; export const CURRENT_SESSION_VERSION = 2; @@ -59,9 +66,12 @@ export interface CompactionEntry extends SessionEntryBase { details?: T; } -export interface BranchSummaryEntry extends SessionEntryBase { +export interface BranchSummaryEntry extends SessionEntryBase { type: "branch_summary"; + fromId: string; summary: string; + /** Hook-specific data (not sent to LLM) */ + details?: T; } /** @@ -145,35 +155,6 @@ export interface SessionInfo { allMessagesText: string; } -export const SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary: - - -`; - -export const SUMMARY_SUFFIX = ` -`; - -/** Exported for compaction.test.ts */ -export function createSummaryMessage(summary: string, timestamp: string): AgentMessage { - return { - role: "user", - content: SUMMARY_PREFIX + summary + SUMMARY_SUFFIX, - timestamp: new Date(timestamp).getTime(), - }; -} - -/** Convert CustomMessageEntry to AgentMessage format */ -function createCustomMessage(entry: CustomMessageEntry): AgentMessage { - return { - role: "hookMessage", - customType: entry.customType, - content: entry.content, - display: entry.display, - details: entry.details, - timestamp: new Date(entry.timestamp).getTime(), - } as AgentMessage; -} - /** Generate a unique short ID (8 hex chars, collision-checked) */ function generateId(byId: { has(id: string): boolean }): string { for (let i = 0; i < 100; i++) { @@ -328,9 +309,21 @@ export function buildSessionContext( // 3. Emit messages after compaction const messages: AgentMessage[] = []; + const appendMessage = (entry: SessionEntry) => { + if (entry.type === "message") { + messages.push(entry.message); + } else if (entry.type === "custom_message") { + messages.push( + createHookMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp), + ); + } else if (entry.type === "branch_summary" && entry.summary) { + messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)); + } + }; + if (compaction) { // Emit summary first - messages.push(createSummaryMessage(compaction.summary, compaction.timestamp)); + messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp)); // Find compaction index in path const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id); @@ -343,37 +336,19 @@ export function buildSessionContext( foundFirstKept = true; } if (foundFirstKept) { - if (entry.type === "message") { - messages.push(entry.message); - } else if (entry.type === "custom_message") { - messages.push(createCustomMessage(entry)); - } else if (entry.type === "branch_summary") { - messages.push(createSummaryMessage(entry.summary, entry.timestamp)); - } + appendMessage(entry); } } // Emit messages after compaction for (let i = compactionIdx + 1; i < path.length; i++) { const entry = path[i]; - if (entry.type === "message") { - messages.push(entry.message); - } else if (entry.type === "custom_message") { - messages.push(createCustomMessage(entry)); - } else if (entry.type === "branch_summary") { - messages.push(createSummaryMessage(entry.summary, entry.timestamp)); - } + appendMessage(entry); } } else { // No compaction - emit all messages, handle branch summaries and custom messages for (const entry of path) { - if (entry.type === "message") { - messages.push(entry.message); - } else if (entry.type === "custom_message") { - messages.push(createCustomMessage(entry)); - } else if (entry.type === "branch_summary") { - messages.push(createSummaryMessage(entry.summary, entry.timestamp)); - } + appendMessage(entry); } } @@ -597,8 +572,13 @@ export class SessionManager { this._persist(entry); } - /** Append a message as child of current leaf, then advance leaf. Returns entry id. */ - appendMessage(message: AgentMessage): string { + /** Append a message as child of current leaf, then advance leaf. Returns entry id. + * Does not allow writing CompactionSummaryMessage and BranchSummaryMessage directly. + * Reason: we want these to be top-level entries in the session, not message session entries, + * so it is easier to find them. + * These need to be appended via appendCompaction() and appendBranchSummary() methods. + */ + appendMessage(message: Message | HookMessage | BashExecutionMessage): string { const entry: SessionMessageEntry = { type: "message", id: generateId(this.byId), @@ -851,6 +831,7 @@ export class SessionManager { id: generateId(this.byId), parentId: branchFromId, timestamp: new Date().toISOString(), + fromId: branchFromId, summary, }; this._appendEntry(entry); diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index d33b8ca9..01c0d689 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -113,7 +113,6 @@ export { CURRENT_SESSION_VERSION, type CustomEntry, type CustomMessageEntry, - createSummaryMessage, type FileEntry, getLatestCompactionEntry, type ModelChangeEntry, @@ -126,8 +125,6 @@ export { type SessionInfo, SessionManager, type SessionMessageEntry, - SUMMARY_PREFIX, - SUMMARY_SUFFIX, type ThinkingLevelChangeEntry, } from "./core/session-manager.js"; export { diff --git a/packages/coding-agent/src/modes/interactive/components/compaction.ts b/packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts similarity index 63% rename from packages/coding-agent/src/modes/interactive/components/compaction.ts rename to packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts index f2835ee7..049880a4 100644 --- a/packages/coding-agent/src/modes/interactive/components/compaction.ts +++ b/packages/coding-agent/src/modes/interactive/components/compaction-summary-message.ts @@ -1,20 +1,19 @@ import { Container, Markdown, Spacer, Text } from "@mariozechner/pi-tui"; +import type { CompactionSummaryMessage } from "packages/coding-agent/src/core/messages.js"; import { getMarkdownTheme, theme } from "../theme/theme.js"; /** - * Component that renders a compaction indicator with collapsed/expanded state. + * Component that renders a compaction message with collapsed/expanded state. * Collapsed: shows "Context compacted from X tokens" * Expanded: shows the full summary rendered as markdown (like a user message) */ -export class CompactionComponent extends Container { +export class CompactionSummaryMessageComponent extends Container { private expanded = false; - private tokensBefore: number; - private summary: string; + private message: CompactionSummaryMessage; - constructor(tokensBefore: number, summary: string) { + constructor(message: CompactionSummaryMessage) { super(); - this.tokensBefore = tokensBefore; - this.summary = summary; + this.message = message; this.updateDisplay(); } @@ -29,9 +28,9 @@ export class CompactionComponent extends Container { if (this.expanded) { // Show header + summary as markdown (like user message) this.addChild(new Spacer(1)); - const header = `**Context compacted from ${this.tokensBefore.toLocaleString()} tokens**\n\n`; + const header = `**Context compacted from ${this.message.tokensBefore.toLocaleString()} tokens**\n\n`; this.addChild( - new Markdown(header + this.summary, 1, 1, getMarkdownTheme(), { + new Markdown(header + this.message.summary, 1, 1, getMarkdownTheme(), { bgColor: (text: string) => theme.bg("userMessageBg", text), color: (text: string) => theme.fg("userMessageText", text), }), @@ -39,7 +38,7 @@ export class CompactionComponent extends Container { this.addChild(new Spacer(1)); } else { // Collapsed: simple text in warning color with token count - const tokenStr = this.tokensBefore.toLocaleString(); + const tokenStr = this.message.tokensBefore.toLocaleString(); this.addChild( new Text( theme.fg("warning", `Earlier messages compacted from ${tokenStr} tokens (ctrl+o to expand)`), diff --git a/packages/coding-agent/src/modes/interactive/components/user-message.ts b/packages/coding-agent/src/modes/interactive/components/user-message.ts index dfeee875..8b95a3b2 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message.ts @@ -5,13 +5,9 @@ import { getMarkdownTheme, theme } from "../theme/theme.js"; * Component that renders a user message */ export class UserMessageComponent extends Container { - constructor(text: string, isFirst: boolean) { + constructor(text: string) { super(); - - // Add spacer before user message (except first one) - if (!isFirst) { - this.addChild(new Spacer(1)); - } + this.addChild(new Spacer(1)); this.addChild( new Markdown(text, 1, 1, getMarkdownTheme(), { bgColor: (text: string) => theme.bg("userMessageBg", text), diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 5b6e6844..405bcf23 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -28,14 +28,8 @@ import { APP_NAME, getAuthPath, getDebugLogPath } from "../../config.js"; import type { AgentSession, AgentSessionEvent } from "../../core/agent-session.js"; import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "../../core/custom-tools/index.js"; import type { HookUIContext } from "../../core/hooks/index.js"; -import { isBashExecutionMessage, isHookMessage } from "../../core/messages.js"; -import { - getLatestCompactionEntry, - type SessionContext, - SessionManager, - SUMMARY_PREFIX, - SUMMARY_SUFFIX, -} from "../../core/session-manager.js"; +import { createCompactionSummaryMessage } from "../../core/messages.js"; +import { type SessionContext, SessionManager } from "../../core/session-manager.js"; import { loadSkills } from "../../core/skills.js"; import { loadProjectContextFiles } from "../../core/system-prompt.js"; import type { TruncationResult } from "../../core/tools/truncate.js"; @@ -44,7 +38,7 @@ import { copyToClipboard } from "../../utils/clipboard.js"; import { ArminComponent } from "./components/armin.js"; import { AssistantMessageComponent } from "./components/assistant-message.js"; import { BashExecutionComponent } from "./components/bash-execution.js"; -import { CompactionComponent } from "./components/compaction.js"; +import { CompactionSummaryMessageComponent } from "./components/compaction-summary-message.js"; import { CustomEditor } from "./components/custom-editor.js"; import { DynamicBorder } from "./components/dynamic-border.js"; import { FooterComponent } from "./components/footer.js"; @@ -84,9 +78,6 @@ export class InteractiveMode { // Tool execution tracking: toolCallId -> component private pendingTools = new Map(); - // Track if this is the first user message (to skip spacer) - private isFirstUserMessage = true; - // Tool output expansion state private toolOutputExpanded = false; @@ -817,7 +808,7 @@ export class InteractiveMode { break; case "message_start": - if (isHookMessage(event.message)) { + if (event.message.role === "hookMessage") { this.addMessageToChat(event.message); this.ui.requestRender(); } else if (event.message.role === "user") { @@ -828,7 +819,7 @@ export class InteractiveMode { } else if (event.message.role === "assistant") { this.streamingComponent = new AssistantMessageComponent(undefined, this.hideThinkingBlock); this.chatContainer.addChild(this.streamingComponent); - this.streamingComponent.updateContent(event.message as AssistantMessage); + this.streamingComponent.updateContent(event.message); this.ui.requestRender(); } break; @@ -983,7 +974,12 @@ export class InteractiveMode { this.chatContainer.clear(); this.rebuildChatFromMessages(); // Add compaction component (same as manual /compact) - const compactionComponent = new CompactionComponent(event.result.tokensBefore, event.result.summary); + const compactionComponent = new CompactionSummaryMessageComponent({ + role: "compactionSummary", + tokensBefore: event.result.tokensBefore, + summary: event.result.summary, + timestamp: Date.now(), + }); compactionComponent.setExpanded(this.toolOutputExpanded); this.chatContainer.addChild(compactionComponent); this.footer.updateState(this.session.state); @@ -1051,38 +1047,70 @@ export class InteractiveMode { this.ui.requestRender(); } - private addMessageToChat(message: AgentMessage): void { - if (isBashExecutionMessage(message)) { - const component = new BashExecutionComponent(message.command, this.ui); - if (message.output) { - component.appendOutput(message.output); + private addMessageToChat(message: AgentMessage, options?: { populateHistory?: boolean }): void { + switch (message.role) { + case "bashExecution": { + const component = new BashExecutionComponent(message.command, this.ui); + if (message.output) { + component.appendOutput(message.output); + } + component.setComplete( + message.exitCode, + message.cancelled, + message.truncated ? ({ truncated: true } as TruncationResult) : undefined, + message.fullOutputPath, + ); + this.chatContainer.addChild(component); + break; } - component.setComplete( - message.exitCode, - message.cancelled, - message.truncated ? ({ truncated: true } as TruncationResult) : undefined, - message.fullOutputPath, - ); - this.chatContainer.addChild(component); - return; - } - - if (isHookMessage(message)) { - // Render as custom message if display is true - if (message.display) { - const renderer = this.session.hookRunner?.getMessageRenderer(message.customType); - this.chatContainer.addChild(new HookMessageComponent(message, renderer)); + case "hookMessage": { + if (message.display) { + const renderer = this.session.hookRunner?.getMessageRenderer(message.customType); + this.chatContainer.addChild(new HookMessageComponent(message, renderer)); + } + break; } - } else if (message.role === "user") { - const textContent = this.getUserMessageText(message); - if (textContent) { - const userComponent = new UserMessageComponent(textContent, this.isFirstUserMessage); - this.chatContainer.addChild(userComponent); - this.isFirstUserMessage = false; + case "compactionSummary": { + const component = new CompactionSummaryMessageComponent(message); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + break; + } + case "branchSummary": { + // Branch summaries are rendered as compaction summaries + const component = new CompactionSummaryMessageComponent({ + role: "compactionSummary", + summary: message.summary, + tokensBefore: 0, + timestamp: message.timestamp, + }); + component.setExpanded(this.toolOutputExpanded); + this.chatContainer.addChild(component); + break; + } + case "user": { + const textContent = this.getUserMessageText(message); + if (textContent) { + const userComponent = new UserMessageComponent(textContent); + this.chatContainer.addChild(userComponent); + if (options?.populateHistory) { + this.editor.addToHistory(textContent); + } + } + break; + } + case "assistant": { + const assistantComponent = new AssistantMessageComponent(message, this.hideThinkingBlock); + this.chatContainer.addChild(assistantComponent); + break; + } + case "toolResult": { + // Tool results are rendered inline with tool calls, handled separately + break; + } + default: { + const _exhaustive: never = message; } - } else if (message.role === "assistant") { - const assistantComponent = new AssistantMessageComponent(message as AssistantMessage, this.hideThinkingBlock); - this.chatContainer.addChild(assistantComponent); } } @@ -1096,7 +1124,6 @@ export class InteractiveMode { sessionContext: SessionContext, options: { updateFooter?: boolean; populateHistory?: boolean } = {}, ): void { - this.isFirstUserMessage = true; this.pendingTools.clear(); if (options.updateFooter) { @@ -1104,65 +1131,25 @@ export class InteractiveMode { this.updateEditorBorderColor(); } - const compactionEntry = getLatestCompactionEntry(this.sessionManager.getEntries()); - - for (let i = 0; i < sessionContext.messages.length; i++) { - const message = sessionContext.messages[i]; - - if (isBashExecutionMessage(message)) { + for (const message of sessionContext.messages) { + // Assistant messages need special handling for tool calls + if (message.role === "assistant") { this.addMessageToChat(message); - continue; - } - - // Check if this is a custom_message entry - if (isHookMessage(message)) { - if (message.display) { - const renderer = this.session.hookRunner?.getMessageRenderer(message.customType); - this.chatContainer.addChild(new HookMessageComponent(message, renderer)); - } - continue; - } - - if (message.role === "user") { - const textContent = this.getUserMessageText(message); - if (textContent) { - if (textContent.startsWith(SUMMARY_PREFIX) && compactionEntry) { - const summary = textContent.slice(SUMMARY_PREFIX.length, -SUMMARY_SUFFIX.length); - const component = new CompactionComponent(compactionEntry.tokensBefore, summary); - component.setExpanded(this.toolOutputExpanded); - this.chatContainer.addChild(component); - } else { - const userComponent = new UserMessageComponent(textContent, this.isFirstUserMessage); - this.chatContainer.addChild(userComponent); - this.isFirstUserMessage = false; - if (options.populateHistory) { - this.editor.addToHistory(textContent); - } - } - } - } else if (message.role === "assistant") { - const assistantMsg = message as AssistantMessage; - const assistantComponent = new AssistantMessageComponent(assistantMsg, this.hideThinkingBlock); - this.chatContainer.addChild(assistantComponent); - - for (const content of assistantMsg.content) { + // Render tool call components + for (const content of message.content) { if (content.type === "toolCall") { const component = new ToolExecutionComponent( content.name, content.arguments, - { - showImages: this.settingsManager.getShowImages(), - }, + { showImages: this.settingsManager.getShowImages() }, this.customTools.get(content.name)?.tool, this.ui, ); this.chatContainer.addChild(component); - if (assistantMsg.stopReason === "aborted" || assistantMsg.stopReason === "error") { + if (message.stopReason === "aborted" || message.stopReason === "error") { const errorMessage = - assistantMsg.stopReason === "aborted" - ? "Operation aborted" - : assistantMsg.errorMessage || "Error"; + message.stopReason === "aborted" ? "Operation aborted" : message.errorMessage || "Error"; component.updateResult({ content: [{ type: "text", text: errorMessage }], isError: true }); } else { this.pendingTools.set(content.id, component); @@ -1170,13 +1157,18 @@ export class InteractiveMode { } } } else if (message.role === "toolResult") { + // Match tool results to pending tool components const component = this.pendingTools.get(message.toolCallId); if (component) { component.updateResult(message); this.pendingTools.delete(message.toolCallId); } + } else { + // All other messages use standard rendering + this.addMessageToChat(message, options); } } + this.pendingTools.clear(); this.ui.requestRender(); } @@ -1308,7 +1300,7 @@ export class InteractiveMode { for (const child of this.chatContainer.children) { if (child instanceof ToolExecutionComponent) { child.setExpanded(this.toolOutputExpanded); - } else if (child instanceof CompactionComponent) { + } else if (child instanceof CompactionSummaryMessageComponent) { child.setExpanded(this.toolOutputExpanded); } else if (child instanceof BashExecutionComponent) { child.setExpanded(this.toolOutputExpanded); @@ -1584,7 +1576,6 @@ export class InteractiveMode { } this.chatContainer.clear(); - this.isFirstUserMessage = true; this.renderInitialMessages(); this.editor.setText(result.selectedText); done(); @@ -1638,7 +1629,6 @@ export class InteractiveMode { // Clear and re-render the chat this.chatContainer.clear(); - this.isFirstUserMessage = true; this.renderInitialMessages(); this.showStatus("Resumed session"); } @@ -1899,7 +1889,6 @@ export class InteractiveMode { this.pendingMessagesContainer.clear(); this.streamingComponent = null; this.pendingTools.clear(); - this.isFirstUserMessage = true; this.chatContainer.addChild(new Spacer(1)); this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1)); @@ -2027,11 +2016,11 @@ export class InteractiveMode { const result = await this.session.compact(customInstructions); // Rebuild UI - this.chatContainer.clear(); this.rebuildChatFromMessages(); // Add compaction component - const compactionComponent = new CompactionComponent(result.tokensBefore, result.summary); + const msg = createCompactionSummaryMessage(result.summary, result.tokensBefore, new Date().toISOString()); + const compactionComponent = new CompactionSummaryMessageComponent(msg); compactionComponent.setExpanded(this.toolOutputExpanded); this.chatContainer.addChild(compactionComponent); diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 84ce3b3c..43f56712 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -16,7 +16,6 @@ import { import { buildSessionContext, type CompactionEntry, - createSummaryMessage, type ModelChangeEntry, migrateSessionEntries, parseSessionEntries, @@ -272,21 +271,6 @@ describe("findCutPoint", () => { }); }); -describe("createSummaryMessage", () => { - it("should create user message with prefix and correct timestamp", () => { - const ts = "2025-01-01T12:00:00.000Z"; - const msg = createSummaryMessage("This is the summary", ts); - expect(msg.role).toBe("user"); - expect(msg.timestamp).toBe(new Date(ts).getTime()); - if (msg.role === "user") { - expect(msg.content).toContain( - "The conversation history before this point was compacted into the following summary:", - ); - expect(msg.content).toContain("This is the summary"); - } - }); -}); - describe("buildSessionContext", () => { it("should load all messages when no compaction", () => { const entries: SessionEntry[] = [ diff --git a/packages/coding-agent/test/session-manager/build-context.test.ts b/packages/coding-agent/test/session-manager/build-context.test.ts index cb34bb84..31479b7b 100644 --- a/packages/coding-agent/test/session-manager/build-context.test.ts +++ b/packages/coding-agent/test/session-manager/build-context.test.ts @@ -6,7 +6,6 @@ import { type ModelChangeEntry, type SessionEntry, type SessionMessageEntry, - SUMMARY_PREFIX, type ThinkingLevelChangeEntry, } from "../../src/core/session-manager.js"; @@ -49,8 +48,8 @@ function compaction(id: string, parentId: string | null, summary: string, firstK }; } -function branchSummary(id: string, parentId: string | null, summary: string): BranchSummaryEntry { - return { type: "branch_summary", id, parentId, timestamp: "2025-01-01T00:00:00Z", summary }; +function branchSummary(id: string, parentId: string | null, fromId: string, summary: string): BranchSummaryEntry { + return { type: "branch_summary", id, parentId, timestamp: "2025-01-01T00:00:00Z", summary, fromId }; } function thinkingLevel(id: string, parentId: string | null, level: string): ThinkingLevelChangeEntry { @@ -151,7 +150,7 @@ describe("buildSessionContext", () => { // Summary + all messages (1,2,4) expect(ctx.messages).toHaveLength(4); - expect((ctx.messages[0] as any).content).toContain(SUMMARY_PREFIX); + expect((ctx.messages[0] as any).content).toContain("Empty summary"); }); it("multiple compactions uses latest", () => { @@ -198,7 +197,7 @@ describe("buildSessionContext", () => { msg("1", null, "user", "start"), msg("2", "1", "assistant", "response"), msg("3", "2", "user", "abandoned path"), - branchSummary("4", "2", "Summary of abandoned work"), + branchSummary("4", "2", "Summary of abandoned work", "3"), msg("5", "4", "user", "new direction"), ]; const ctx = buildSessionContext(entries, "5"); @@ -225,7 +224,7 @@ describe("buildSessionContext", () => { msg("8", "3", "user", "wrong path"), msg("9", "8", "assistant", "wrong response"), // Branch summary resuming from 3 - branchSummary("10", "3", "Tried wrong approach"), + branchSummary("10", "3", "Tried wrong approach", "9"), msg("11", "10", "user", "better approach"), ];