From 159075cad7aff8b6e861c14140babb545492d134 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 11 Nov 2025 23:05:58 +0100 Subject: [PATCH] Improve tool execution rendering and error handling - Show tool execution components immediately when tool calls appear in streaming - Update components with streaming arguments as they come in - Handle incomplete/partial arguments gracefully with optional chaining - Fix error handling: tools now throw exceptions instead of returning error messages - Fix bash abort handling to properly reject on abort/timeout - Clean up error display --- packages/coding-agent/example.txt | 57 +++++- packages/coding-agent/src/tools/bash.ts | 19 +- packages/coding-agent/src/tools/edit.ts | 80 ++++---- packages/coding-agent/src/tools/read.ts | 30 ++- packages/coding-agent/src/tools/write.ts | 30 ++- .../coding-agent/src/tui/assistant-message.ts | 56 +----- .../coding-agent/src/tui/tool-execution.ts | 112 ++++++++---- packages/coding-agent/src/tui/tui-renderer.ts | 172 ++++++------------ packages/coding-agent/src/tui/user-message.ts | 4 +- packages/tui/src/components/loader.ts | 6 +- 10 files changed, 288 insertions(+), 278 deletions(-) diff --git a/packages/coding-agent/example.txt b/packages/coding-agent/example.txt index db239d22..15f02622 100644 --- a/packages/coding-agent/example.txt +++ b/packages/coding-agent/example.txt @@ -1,14 +1,14 @@ The Amazing Adventures of Fox and Dog ====================================== -Long ago, in a magical forest clearing, there lived an extraordinarily swift brown fox. -This legendary fox was famous across the entire woodland for its blazing speed and grace. +Long ago, in a mystical forest clearing, there lived an incredibly fast brown fox. +This legendary fox was renowned throughout the entire woodland for its incredible speed and agility. Each dawn, the fox would sprint through the ancient trees, soaring over logs and babbling brooks. The woodland creatures gazed in wonder as it flashed past them like a streak of copper lightning. -At the clearing's edge, there also dwelled an exceptionally lazy dog. -This contented dog much preferred snoozing in the golden sunshine to any kind of adventure. +At the clearing's edge, there also lived a very lazy dog. +This happy dog much preferred napping in the warm sunshine to any kind of adventure. One fateful morning, the fox challenged the dog to an epic race across the meadow. The dog yawned deeply and declined, saying "Why rush around when you can rest peacefully?" @@ -27,4 +27,53 @@ Other times, stillness and contentment are the greatest treasures. Both the quick fox and the lazy dog lived happily in their own ways. +And so their friendship grew stronger with each passing season. +The fox would return from adventures with tales of distant lands. +The dog would listen contentedly, never needing to leave home. +They learned that happiness comes in many different forms. +The forest creatures admired their unlikely bond. +Some days the fox would rest beside the dog in the sunshine. +Other days the dog would take a short stroll with the fox. +They discovered balance between motion and stillness. +The wise old owl observed them from his towering oak. +He noted that both had found their true nature. +Winter came and blanketed the forest in sparkling snow. +The fox's copper fur stood out against the white landscape. +The dog found a cozy spot by the warmest rock. +They shared stories as snowflakes drifted down around them. +Spring arrived with flowers blooming across the meadow. +The fox chased butterflies through fields of wildflowers. +The dog rolled in patches of soft clover and sweet grass. +Summer brought long days of golden light and warmth. +The fox discovered hidden streams in the deep forest. +The dog found the perfect shady spot beneath an elm tree. +Autumn painted the woods in brilliant reds and golds. +The fox leaped through piles of crunchy fallen leaves. +The dog watched the changing colors from his favorite perch. +Years passed and both grew wiser in their own ways. +The fox learned when to rest and the dog learned when to play. +Young animals would visit to hear their wisdom. +"Be true to yourself," the fox would always say. +"Find joy in your own path," the dog would add. +Their story spread throughout the woodland realm. +It became a tale told to every new generation. +Parents would share it with their curious young ones. +Teachers would use it in lessons about acceptance. +Travelers would stop to see the famous pair. +Artists painted pictures of the fox and dog together. +Poets wrote verses about their enduring friendship. +Musicians composed songs celebrating their harmony. +The clearing became a place of peace and understanding. +All creatures were welcome to rest there. +The fox still runs when the spirit moves him. +The dog still naps when the mood strikes him. +Neither judges the other for their choices. +Both have found contentment in being themselves. +The moon rises over the peaceful forest each night. +Stars twinkle above the quiet clearing. +The fox and dog sleep side by side. +Dreams of adventure and rest mingle together. +Morning will bring new possibilities for both. +But tonight, all is calm and perfect. +This is how true friendship looks. The End. diff --git a/packages/coding-agent/src/tools/bash.ts b/packages/coding-agent/src/tools/bash.ts index a3e6630f..3116d12d 100644 --- a/packages/coding-agent/src/tools/bash.ts +++ b/packages/coding-agent/src/tools/bash.ts @@ -13,7 +13,7 @@ export const bashTool: AgentTool = { "Execute a bash command in the current working directory. Returns stdout and stderr. Commands run with a 30 second timeout.", parameters: bashSchema, execute: async (_toolCallId: string, { command }: { command: string }, signal?: AbortSignal) => { - return new Promise((resolve) => { + return new Promise((resolve, reject) => { const child = spawn("sh", ["-c", command], { detached: true, stdio: ["ignore", "pipe", "pipe"], @@ -65,10 +65,9 @@ export const bashTool: AgentTool = { if (output) output += "\n"; output += stderr; } - resolve({ - output: output || "(no output)", - details: undefined, - }); + if (output) output += "\n\n"; + output += "Command aborted"; + reject(new Error(output)); return; } @@ -81,10 +80,7 @@ export const bashTool: AgentTool = { } if (output) output += "\n\n"; output += "Command timed out after 30 seconds"; - resolve({ - output, - details: undefined, - }); + reject(new Error(output)); return; } @@ -97,10 +93,7 @@ export const bashTool: AgentTool = { if (code !== 0 && code !== null) { if (output) output += "\n\n"; - resolve({ - output: `${output}Command exited with code ${code}`, - details: undefined, - }); + reject(new Error(`${output}Command exited with code ${code}`)); } else { resolve({ output: output || "(no output)", details: undefined }); } diff --git a/packages/coding-agent/src/tools/edit.ts b/packages/coding-agent/src/tools/edit.ts index 57495533..e1ee2005 100644 --- a/packages/coding-agent/src/tools/edit.ts +++ b/packages/coding-agent/src/tools/edit.ts @@ -1,8 +1,22 @@ +import * as os from "node:os"; import type { AgentTool } from "@mariozechner/pi-ai"; import { Type } from "@sinclair/typebox"; import { existsSync, readFileSync, writeFileSync } from "fs"; import { resolve } from "path"; +/** + * Expand ~ to home directory + */ +function expandPath(filePath: string): string { + if (filePath === "~") { + return os.homedir(); + } + if (filePath.startsWith("~/")) { + return os.homedir() + filePath.slice(1); + } + return filePath; +} + const editSchema = Type.Object({ path: Type.String({ description: "Path to the file to edit (relative or absolute)" }), oldText: Type.String({ description: "Exact text to find and replace (must match exactly)" }), @@ -19,43 +33,37 @@ export const editTool: AgentTool = { _toolCallId: string, { path, oldText, newText }: { path: string; oldText: string; newText: string }, ) => { - try { - const absolutePath = resolve(path); + const absolutePath = resolve(expandPath(path)); - if (!existsSync(absolutePath)) { - return { output: `Error: File not found: ${path}`, details: undefined }; - } - - const content = readFileSync(absolutePath, "utf-8"); - - // Check if old text exists - if (!content.includes(oldText)) { - return { - output: `Error: Could not find the exact text in ${path}. The old text must match exactly including all whitespace and newlines.`, - details: undefined, - }; - } - - // Count occurrences - const occurrences = content.split(oldText).length - 1; - - if (occurrences > 1) { - return { - output: `Error: Found ${occurrences} occurrences of the text in ${path}. The text must be unique. Please provide more context to make it unique.`, - details: undefined, - }; - } - - // Perform replacement - const newContent = content.replace(oldText, newText); - writeFileSync(absolutePath, newContent, "utf-8"); - - return { - output: `Successfully replaced text in ${path}. Changed ${oldText.length} characters to ${newText.length} characters.`, - details: undefined, - }; - } catch (error: any) { - return { output: `Error editing file: ${error.message}`, details: undefined }; + if (!existsSync(absolutePath)) { + throw new Error(`File not found: ${path}`); } + + const content = readFileSync(absolutePath, "utf-8"); + + // Check if old text exists + if (!content.includes(oldText)) { + throw new Error( + `Could not find the exact text in ${path}. The old text must match exactly including all whitespace and newlines.`, + ); + } + + // Count occurrences + const occurrences = content.split(oldText).length - 1; + + if (occurrences > 1) { + throw new Error( + `Found ${occurrences} occurrences of the text in ${path}. The text must be unique. Please provide more context to make it unique.`, + ); + } + + // Perform replacement + const newContent = content.replace(oldText, newText); + writeFileSync(absolutePath, newContent, "utf-8"); + + return { + output: `Successfully replaced text in ${path}. Changed ${oldText.length} characters to ${newText.length} characters.`, + details: undefined, + }; }, }; diff --git a/packages/coding-agent/src/tools/read.ts b/packages/coding-agent/src/tools/read.ts index bde7cc95..9e549cda 100644 --- a/packages/coding-agent/src/tools/read.ts +++ b/packages/coding-agent/src/tools/read.ts @@ -1,8 +1,22 @@ +import * as os from "node:os"; import type { AgentTool } from "@mariozechner/pi-ai"; import { Type } from "@sinclair/typebox"; import { existsSync, readFileSync } from "fs"; import { resolve } from "path"; +/** + * Expand ~ to home directory + */ +function expandPath(filePath: string): string { + if (filePath === "~") { + return os.homedir(); + } + if (filePath.startsWith("~/")) { + return os.homedir() + filePath.slice(1); + } + return filePath; +} + const readSchema = Type.Object({ path: Type.String({ description: "Path to the file to read (relative or absolute)" }), }); @@ -13,17 +27,13 @@ export const readTool: AgentTool = { description: "Read the contents of a file. Returns the full file content as text.", parameters: readSchema, execute: async (_toolCallId: string, { path }: { path: string }) => { - try { - const absolutePath = resolve(path); + const absolutePath = resolve(expandPath(path)); - if (!existsSync(absolutePath)) { - return { output: `Error: File not found: ${path}`, details: undefined }; - } - - const content = readFileSync(absolutePath, "utf-8"); - return { output: content, details: undefined }; - } catch (error: any) { - return { output: `Error reading file: ${error.message}`, details: undefined }; + if (!existsSync(absolutePath)) { + throw new Error(`File not found: ${path}`); } + + const content = readFileSync(absolutePath, "utf-8"); + return { output: content, details: undefined }; }, }; diff --git a/packages/coding-agent/src/tools/write.ts b/packages/coding-agent/src/tools/write.ts index ef8b2d5e..4eb18e25 100644 --- a/packages/coding-agent/src/tools/write.ts +++ b/packages/coding-agent/src/tools/write.ts @@ -1,8 +1,22 @@ +import * as os from "node:os"; import type { AgentTool } from "@mariozechner/pi-ai"; import { Type } from "@sinclair/typebox"; import { mkdirSync, writeFileSync } from "fs"; import { dirname, resolve } from "path"; +/** + * Expand ~ to home directory + */ +function expandPath(filePath: string): string { + if (filePath === "~") { + return os.homedir(); + } + if (filePath.startsWith("~/")) { + return os.homedir() + filePath.slice(1); + } + return filePath; +} + const writeSchema = Type.Object({ path: Type.String({ description: "Path to the file to write (relative or absolute)" }), content: Type.String({ description: "Content to write to the file" }), @@ -15,17 +29,13 @@ export const writeTool: AgentTool = { "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Automatically creates parent directories.", parameters: writeSchema, execute: async (_toolCallId: string, { path, content }: { path: string; content: string }) => { - try { - const absolutePath = resolve(path); - const dir = dirname(absolutePath); + const absolutePath = resolve(expandPath(path)); + const dir = dirname(absolutePath); - // Create parent directories if needed - mkdirSync(dir, { recursive: true }); + // Create parent directories if needed + mkdirSync(dir, { recursive: true }); - writeFileSync(absolutePath, content, "utf-8"); - return { output: `Successfully wrote ${content.length} bytes to ${path}`, details: undefined }; - } catch (error: any) { - return { output: `Error writing file: ${error.message}`, details: undefined }; - } + writeFileSync(absolutePath, content, "utf-8"); + return { output: `Successfully wrote ${content.length} bytes to ${path}`, details: undefined }; }, }; diff --git a/packages/coding-agent/src/tui/assistant-message.ts b/packages/coding-agent/src/tui/assistant-message.ts index b245e90e..95529012 100644 --- a/packages/coding-agent/src/tui/assistant-message.ts +++ b/packages/coding-agent/src/tui/assistant-message.ts @@ -6,25 +6,15 @@ import chalk from "chalk"; * Component that renders a complete assistant message */ export class AssistantMessageComponent extends Container { - private spacer: Spacer; private contentContainer: Container; - private statsText: Text; constructor(message?: AssistantMessage) { super(); - // Add spacer before assistant message - this.spacer = new Spacer(1); - this.addChild(this.spacer); - // Container for text/thinking content this.contentContainer = new Container(); this.addChild(this.contentContainer); - // Stats text - this.statsText = new Text("", 1, 0); - this.addChild(this.statsText); - if (message) { this.updateContent(message); } @@ -34,8 +24,14 @@ export class AssistantMessageComponent extends Container { // Clear content container this.contentContainer.clear(); - // Check if there's any actual content (text or thinking) - let hasContent = false; + if ( + message.content.length > 0 && + message.content.some( + (c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim()), + ) + ) { + this.contentContainer.addChild(new Spacer(1)); + } // Render content in order for (const content of message.content) { @@ -43,57 +39,21 @@ export class AssistantMessageComponent extends Container { // Assistant text messages with no background - trim the text // Set paddingY=0 to avoid extra spacing before tool executions this.contentContainer.addChild(new Markdown(content.text.trim(), undefined, undefined, undefined, 1, 0)); - hasContent = true; } else if (content.type === "thinking" && content.thinking.trim()) { // Thinking traces in dark gray italic // Apply styling to entire block so wrapping preserves it const thinkingText = chalk.gray.italic(content.thinking); this.contentContainer.addChild(new Text(thinkingText, 1, 0)); this.contentContainer.addChild(new Spacer(1)); - hasContent = true; } } // Check if aborted - show after partial content if (message.stopReason === "aborted") { this.contentContainer.addChild(new Text(chalk.red("Aborted"))); - hasContent = true; } else if (message.stopReason === "error") { const errorMsg = message.errorMessage || "Unknown error"; this.contentContainer.addChild(new Text(chalk.red(`Error: ${errorMsg}`))); - hasContent = true; } - - // Only update stats if there's actual content (not just tool calls) - if (hasContent) { - this.updateStats(message.usage); - } - } - - updateStats(usage: any): void { - if (!usage) { - this.statsText.setText(""); - return; - } - - // Format token counts - const formatTokens = (count: number): string => { - if (count < 1000) return count.toString(); - if (count < 10000) return (count / 1000).toFixed(1) + "k"; - return Math.round(count / 1000) + "k"; - }; - - const statsParts = []; - if (usage.input) statsParts.push(`↑${formatTokens(usage.input)}`); - if (usage.output) statsParts.push(`↓${formatTokens(usage.output)}`); - if (usage.cacheRead) statsParts.push(`R${formatTokens(usage.cacheRead)}`); - if (usage.cacheWrite) statsParts.push(`W${formatTokens(usage.cacheWrite)}`); - if (usage.cost?.total) statsParts.push(`$${usage.cost.total.toFixed(3)}`); - - this.statsText.setText(chalk.gray(statsParts.join(" "))); - } - - hideStats(): void { - this.statsText.setText(""); } } diff --git a/packages/coding-agent/src/tui/tool-execution.ts b/packages/coding-agent/src/tui/tool-execution.ts index b25090a5..21fdf51b 100644 --- a/packages/coding-agent/src/tui/tool-execution.ts +++ b/packages/coding-agent/src/tui/tool-execution.ts @@ -1,11 +1,55 @@ +import * as os from "node:os"; import { Container, Spacer, Text } from "@mariozechner/pi-tui"; import chalk from "chalk"; +/** + * Convert absolute path to tilde notation if it's in home directory + */ +function shortenPath(path: string): string { + const home = os.homedir(); + if (path.startsWith(home)) { + return "~" + path.slice(home.length); + } + return path; +} + +/** + * Generate a unified diff between old and new strings with line numbers + */ +function generateDiff(oldStr: string, newStr: string): string { + // Split into lines + const oldLines = oldStr.split("\n"); + const newLines = newStr.split("\n"); + + const diff: string[] = []; + + // Calculate line number padding (for alignment) + const maxLineNum = Math.max(oldLines.length, newLines.length); + const lineNumWidth = String(maxLineNum).length; + + // Show old lines with line numbers + diff.push(chalk.red("- old:")); + for (let i = 0; i < oldLines.length; i++) { + const lineNum = String(i + 1).padStart(lineNumWidth, " "); + diff.push(chalk.red(`- ${chalk.dim(lineNum)} ${oldLines[i]}`)); + } + + diff.push(""); + + // Show new lines with line numbers + diff.push(chalk.green("+ new:")); + for (let i = 0; i < newLines.length; i++) { + const lineNum = String(i + 1).padStart(lineNumWidth, " "); + diff.push(chalk.green(`+ ${chalk.dim(lineNum)} ${newLines[i]}`)); + } + + return diff.join("\n"); +} + /** * Component that renders a tool call with its result (updateable) */ export class ToolExecutionComponent extends Container { - private spacer: Spacer; private contentText: Text; private toolName: string; private args: any; @@ -15,15 +59,18 @@ export class ToolExecutionComponent extends Container { super(); this.toolName = toolName; this.args = args; - // Blank line with no background for spacing - this.spacer = new Spacer(1); - this.addChild(this.spacer); + this.addChild(new Spacer(1)); // Content with colored background and padding this.contentText = new Text("", 1, 1, { r: 40, g: 40, b: 50 }); this.addChild(this.contentText); this.updateDisplay(); } + updateArgs(args: any): void { + this.args = args; + this.updateDisplay(); + } + updateResult(result: { output: string; isError: boolean }): void { this.result = result; this.updateDisplay(); @@ -45,11 +92,8 @@ export class ToolExecutionComponent extends Container { // Format based on tool type if (this.toolName === "bash") { - const command = this.args.command || ""; - text = chalk.bold(`$ ${command}`); - if (this.result?.isError) { - text += " ❌"; - } + const command = this.args?.command || ""; + text = chalk.bold(`$ ${command || chalk.dim("...")}`); if (this.result) { // Show output without code fences - more minimal @@ -67,11 +111,8 @@ export class ToolExecutionComponent extends Container { } } } else if (this.toolName === "read") { - const path = this.args.path || ""; - text = chalk.bold("read") + " " + chalk.cyan(path); - if (this.result?.isError) { - text += " ❌"; - } + const path = shortenPath(this.args?.file_path || this.args?.path || ""); + text = chalk.bold("read") + " " + (path ? chalk.cyan(path) : chalk.dim("...")); if (this.result) { const lines = this.result.output.split("\n"); @@ -85,43 +126,38 @@ export class ToolExecutionComponent extends Container { } } } else if (this.toolName === "write") { - const path = this.args.path || ""; - const fileContent = this.args.content || ""; - const lines = fileContent.split("\n"); + const path = shortenPath(this.args?.file_path || this.args?.path || ""); + const fileContent = this.args?.content || ""; + const lines = fileContent ? fileContent.split("\n") : []; const totalLines = lines.length; - text = chalk.bold("write") + " " + chalk.cyan(path); + text = chalk.bold("write") + " " + (path ? chalk.cyan(path) : chalk.dim("...")); if (totalLines > 10) { text += ` (${totalLines} lines)`; } - if (this.result) { - text += this.result.isError ? " ❌" : " ✓"; - } + // Show first 10 lines of content if available + if (fileContent) { + const maxLines = 10; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; - // Show first 10 lines of content - const maxLines = 10; - const displayLines = lines.slice(0, maxLines); - const remaining = lines.length - maxLines; - - text += "\n\n" + displayLines.map((line: string) => chalk.dim(line)).join("\n"); - if (remaining > 0) { - text += chalk.dim(`\n... (${remaining} more lines)`); + text += "\n\n" + displayLines.map((line: string) => chalk.dim(line)).join("\n"); + if (remaining > 0) { + text += chalk.dim(`\n... (${remaining} more lines)`); + } } } else if (this.toolName === "edit") { - const path = this.args.path || ""; - text = chalk.bold("edit") + " " + chalk.cyan(path); - if (this.result) { - text += this.result.isError ? " ❌" : " ✓"; + const path = shortenPath(this.args?.file_path || this.args?.path || ""); + text = chalk.bold("edit") + " " + (path ? chalk.cyan(path) : chalk.dim("...")); + + // Show diff if we have old_string and new_string + if (this.args?.old_string && this.args?.new_string) { + text += "\n\n" + generateDiff(this.args.old_string, this.args.new_string); } } else { // Generic tool text = chalk.bold(this.toolName); - if (this.result?.isError) { - text += " ❌"; - } else if (this.result) { - text += " ✓"; - } const content = JSON.stringify(this.args, null, 2); text += "\n\n" + content; diff --git a/packages/coding-agent/src/tui/tui-renderer.ts b/packages/coding-agent/src/tui/tui-renderer.ts index 7ba0b1f8..7ab024fc 100644 --- a/packages/coding-agent/src/tui/tui-renderer.ts +++ b/packages/coding-agent/src/tui/tui-renderer.ts @@ -1,7 +1,15 @@ import type { Agent, AgentEvent, AgentState } from "@mariozechner/pi-agent"; import type { AssistantMessage, Message } from "@mariozechner/pi-ai"; import type { SlashCommand } from "@mariozechner/pi-tui"; -import { CombinedAutocompleteProvider, Container, Loader, ProcessTerminal, Text, TUI } from "@mariozechner/pi-tui"; +import { + CombinedAutocompleteProvider, + Container, + Loader, + ProcessTerminal, + Spacer, + Text, + TUI, +} from "@mariozechner/pi-tui"; import chalk from "chalk"; import { AssistantMessageComponent } from "./assistant-message.js"; import { CustomEditor } from "./custom-editor.js"; @@ -34,9 +42,6 @@ export class TuiRenderer { // Tool execution tracking: toolCallId -> component private pendingTools = new Map(); - // Track assistant message with tool calls that needs stats shown after tools complete - private deferredStats: { component: AssistantMessageComponent; usage: any; toolCallIds: Set } | null = null; - // Thinking level selector private thinkingSelector: ThinkingSelectorComponent | null = null; @@ -88,12 +93,14 @@ export class TuiRenderer { "\n" + chalk.dim("drop files") + chalk.gray(" to attach"); - const header = new Text(logo + "\n" + instructions); + const header = new Text(logo + "\n" + instructions, 1, 0); // Setup UI layout + this.ui.addChild(new Spacer(1)); this.ui.addChild(header); this.ui.addChild(this.chatContainer); this.ui.addChild(this.statusContainer); + this.ui.addChild(new Spacer(1)); this.ui.addChild(this.editorContainer); // Use container that can hold editor or selector this.ui.addChild(this.footer); this.ui.setFocus(this.editor); @@ -173,7 +180,28 @@ export class TuiRenderer { case "message_update": // Update streaming component if (this.streamingComponent && event.message.role === "assistant") { - this.streamingComponent.updateContent(event.message as AssistantMessage); + const assistantMsg = event.message as AssistantMessage; + this.streamingComponent.updateContent(assistantMsg); + + // Create tool execution components as soon as we see tool calls + for (const content of assistantMsg.content) { + if (content.type === "toolCall") { + // Only create if we haven't created it yet + if (!this.pendingTools.has(content.id)) { + this.chatContainer.addChild(new Text("", 0, 0)); + const component = new ToolExecutionComponent(content.name, content.arguments); + this.chatContainer.addChild(component); + this.pendingTools.set(content.id, component); + } else { + // Update existing component with latest arguments as they stream + const component = this.pendingTools.get(content.id); + if (component) { + component.updateArgs(content.arguments); + } + } + } + } + this.ui.requestRender(); } break; @@ -184,23 +212,6 @@ export class TuiRenderer { break; } if (this.streamingComponent && event.message.role === "assistant") { - const assistantMsg = event.message as AssistantMessage; - - // Check if this message has tool calls - const hasToolCalls = assistantMsg.content.some((c) => c.type === "toolCall"); - - if (hasToolCalls) { - // Defer stats until after tool executions complete - const toolCallIds = new Set(); - for (const content of assistantMsg.content) { - if (content.type === "toolCall") { - toolCallIds.add(content.id); - } - } - this.deferredStats = { component: this.streamingComponent, usage: assistantMsg.usage, toolCallIds }; - // Hide stats for now - this.streamingComponent.hideStats(); - } // Keep the streaming component - it's now the final assistant message this.streamingComponent = null; } @@ -208,13 +219,13 @@ export class TuiRenderer { break; case "tool_execution_start": { - // Add empty line before tool execution - this.chatContainer.addChild(new Text("", 0, 0)); - // Create tool execution component and add it - const component = new ToolExecutionComponent(event.toolName, event.args); - this.chatContainer.addChild(component); - this.pendingTools.set(event.toolCallId, component); - this.ui.requestRender(); + // Component should already exist from message_update, but create if missing + if (!this.pendingTools.has(event.toolCallId)) { + const component = new ToolExecutionComponent(event.toolName, event.args); + this.chatContainer.addChild(component); + this.pendingTools.set(event.toolCallId, component); + this.ui.requestRender(); + } break; } @@ -228,17 +239,6 @@ export class TuiRenderer { isError: event.isError, }); this.pendingTools.delete(event.toolCallId); - - // Check if this was part of deferred stats and all tools are complete - if (this.deferredStats) { - this.deferredStats.toolCallIds.delete(event.toolCallId); - if (this.deferredStats.toolCallIds.size === 0) { - // All tools complete - show stats now on the component - this.deferredStats.component.updateStats(this.deferredStats.usage); - this.deferredStats = null; - } - } - this.ui.requestRender(); } break; @@ -256,7 +256,6 @@ export class TuiRenderer { this.streamingComponent = null; } this.pendingTools.clear(); - this.deferredStats = null; // Clear any deferred stats this.editor.disableSubmit = false; this.ui.requestRender(); break; @@ -280,40 +279,12 @@ export class TuiRenderer { // Add assistant message component const assistantComponent = new AssistantMessageComponent(assistantMsg); this.chatContainer.addChild(assistantComponent); - - // Check if this message has tool calls - const hasToolCalls = assistantMsg.content.some((c) => c.type === "toolCall"); - - if (hasToolCalls) { - // Defer stats until after tool executions complete - const toolCallIds = new Set(); - for (const content of assistantMsg.content) { - if (content.type === "toolCall") { - toolCallIds.add(content.id); - } - } - this.deferredStats = { component: assistantComponent, usage: assistantMsg.usage, toolCallIds }; - // Hide stats for now - assistantComponent.hideStats(); - } - // else: stats are shown by the component constructor } // Note: tool calls and results are now handled via tool_execution_start/end events } renderInitialMessages(state: AgentState): void { // Render all existing messages (for --continue mode) - // Track assistant components with their tool calls to show stats after tools - const assistantWithTools = new Map< - number, - { - component: AssistantMessageComponent; - usage: any; - toolCallIds: Set; - remainingToolCallIds: Set; - } - >(); - // Reset first user message flag for initial render this.isFirstUserMessage = true; @@ -335,60 +306,31 @@ export class TuiRenderer { const assistantComponent = new AssistantMessageComponent(assistantMsg); this.chatContainer.addChild(assistantComponent); - // Check if this message has tool calls - const toolCallIds = new Set(); + // Create tool execution components for any tool calls for (const content of assistantMsg.content) { if (content.type === "toolCall") { - toolCallIds.add(content.id); + const component = new ToolExecutionComponent(content.name, content.arguments); + this.chatContainer.addChild(component); + // Store in map so we can update with results later + this.pendingTools.set(content.id, component); } } - if (toolCallIds.size > 0) { - // Hide stats until tools complete - assistantComponent.hideStats(); - assistantWithTools.set(i, { - component: assistantComponent, - usage: assistantMsg.usage, - toolCallIds, - remainingToolCallIds: new Set(toolCallIds), - }); - } } else if (message.role === "toolResult") { - // Render tool calls that have already completed + // Update existing tool execution component with results const toolResultMsg = message as any; - const assistantMsgIndex = state.messages.findIndex( - (m) => - m.role === "assistant" && - m.content.some((c: any) => c.type === "toolCall" && c.id === toolResultMsg.toolCallId), - ); - - if (assistantMsgIndex !== -1) { - const assistantMsg = state.messages[assistantMsgIndex] as AssistantMessage; - const toolCall = assistantMsg.content.find( - (c) => c.type === "toolCall" && c.id === toolResultMsg.toolCallId, - ) as any; - if (toolCall) { - // Add empty line before tool execution - this.chatContainer.addChild(new Text("", 0, 0)); - const component = new ToolExecutionComponent(toolCall.name, toolCall.arguments); - component.updateResult({ - output: toolResultMsg.output, - isError: toolResultMsg.isError, - }); - this.chatContainer.addChild(component); - - // Check if this was the last tool call for this assistant message - const assistantData = assistantWithTools.get(assistantMsgIndex); - if (assistantData) { - assistantData.remainingToolCallIds.delete(toolResultMsg.toolCallId); - if (assistantData.remainingToolCallIds.size === 0) { - // All tools for this assistant message are complete - show stats - assistantData.component.updateStats(assistantData.usage); - } - } - } + const component = this.pendingTools.get(toolResultMsg.toolCallId); + if (component) { + component.updateResult({ + output: toolResultMsg.output, + isError: toolResultMsg.isError, + }); + // Remove from pending map since it's complete + this.pendingTools.delete(toolResultMsg.toolCallId); } } } + // Clear pending tools after rendering initial messages + this.pendingTools.clear(); this.ui.requestRender(); } diff --git a/packages/coding-agent/src/tui/user-message.ts b/packages/coding-agent/src/tui/user-message.ts index d62ab7d4..7369c459 100644 --- a/packages/coding-agent/src/tui/user-message.ts +++ b/packages/coding-agent/src/tui/user-message.ts @@ -4,7 +4,6 @@ import { Container, Markdown, Spacer } from "@mariozechner/pi-tui"; * Component that renders a user message */ export class UserMessageComponent extends Container { - private spacer: Spacer | null = null; private markdown: Markdown; constructor(text: string, isFirst: boolean) { @@ -12,8 +11,7 @@ export class UserMessageComponent extends Container { // Add spacer before user message (except first one) if (!isFirst) { - this.spacer = new Spacer(1); - this.addChild(this.spacer); + this.addChild(new Spacer(1)); } // User messages with dark gray background diff --git a/packages/tui/src/components/loader.ts b/packages/tui/src/components/loader.ts index 39755706..6ea4b84c 100644 --- a/packages/tui/src/components/loader.ts +++ b/packages/tui/src/components/loader.ts @@ -15,11 +15,15 @@ export class Loader extends Text { ui: TUI, private message: string = "Loading...", ) { - super(""); + super("", 1, 0); this.ui = ui; this.start(); } + render(width: number): string[] { + return ["", ...super.render(width)]; + } + start() { this.updateDisplay(); this.intervalId = setInterval(() => {