mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 21:03:42 +00:00
Refactor TUI components into separate files
- Move TUI components into src/tui/ folder - Split out CustomEditor, StreamingMessageComponent, ToolExecutionComponent, FooterComponent - Trim assistant message text content - Add newline after per-message token/cost stats - Improve code organization and maintainability
This commit is contained in:
Mario Zechner
parent
fe5706885d
commit
4fa09814bd
6 changed files with 292 additions and 281 deletions
558
packages/coding-agent/src/tui/tui-renderer.ts
Normal file
558
packages/coding-agent/src/tui/tui-renderer.ts
Normal file
|
|
@ -0,0 +1,558 @@
|
|||
import type { Agent, AgentState, ThinkingLevel } from "@mariozechner/pi-agent";
|
||||
import type { AssistantMessage, Message } from "@mariozechner/pi-ai";
|
||||
import type { SlashCommand } from "@mariozechner/pi-tui";
|
||||
import {
|
||||
CombinedAutocompleteProvider,
|
||||
Container,
|
||||
Loader,
|
||||
Markdown,
|
||||
ProcessTerminal,
|
||||
type SelectItem,
|
||||
SelectList,
|
||||
Text,
|
||||
TUI,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import chalk from "chalk";
|
||||
import { CustomEditor } from "./custom-editor.js";
|
||||
import { FooterComponent } from "./footer.js";
|
||||
import { StreamingMessageComponent } from "./streaming-message.js";
|
||||
import { ToolExecutionComponent } from "./tool-execution.js";
|
||||
|
||||
/**
|
||||
* TUI renderer for the coding agent
|
||||
*/
|
||||
export class TuiRenderer {
|
||||
private ui: TUI;
|
||||
private chatContainer: Container;
|
||||
private statusContainer: Container;
|
||||
private editor: CustomEditor;
|
||||
private editorContainer: Container; // Container to swap between editor and selector
|
||||
private footer: FooterComponent;
|
||||
private agent: Agent;
|
||||
private version: string;
|
||||
private isInitialized = false;
|
||||
private onInputCallback?: (text: string) => void;
|
||||
private loadingAnimation: Loader | null = null;
|
||||
private onInterruptCallback?: () => void;
|
||||
private lastSigintTime = 0;
|
||||
|
||||
// Streaming message tracking
|
||||
private streamingComponent: StreamingMessageComponent | null = null;
|
||||
|
||||
// Tool execution tracking: toolCallId -> component
|
||||
private pendingTools = new Map<string, ToolExecutionComponent>();
|
||||
|
||||
// Track assistant message with tool calls that needs stats shown after tools complete
|
||||
private deferredStats: { usage: any; toolCallIds: Set<string> } | null = null;
|
||||
|
||||
// Thinking level selector
|
||||
private thinkingSelector: SelectList | null = null;
|
||||
|
||||
constructor(agent: Agent, version: string) {
|
||||
this.agent = agent;
|
||||
this.version = version;
|
||||
this.ui = new TUI(new ProcessTerminal());
|
||||
this.chatContainer = new Container();
|
||||
this.statusContainer = new Container();
|
||||
this.editor = new CustomEditor();
|
||||
this.editorContainer = new Container(); // Container to hold editor or selector
|
||||
this.editorContainer.addChild(this.editor); // Start with editor
|
||||
this.footer = new FooterComponent(agent.state);
|
||||
|
||||
// Define slash commands
|
||||
const thinkingCommand: SlashCommand = {
|
||||
name: "thinking",
|
||||
description: "Select reasoning level (opens selector UI)",
|
||||
};
|
||||
|
||||
// Setup autocomplete for file paths and slash commands
|
||||
const autocompleteProvider = new CombinedAutocompleteProvider([thinkingCommand], process.cwd());
|
||||
this.editor.setAutocompleteProvider(autocompleteProvider);
|
||||
}
|
||||
|
||||
async init(): Promise<void> {
|
||||
if (this.isInitialized) return;
|
||||
|
||||
// Add header with logo and instructions
|
||||
const logo = chalk.bold.cyan("pi") + chalk.dim(` v${this.version}`);
|
||||
const instructions =
|
||||
chalk.dim("esc") +
|
||||
chalk.gray(" to interrupt") +
|
||||
"\n" +
|
||||
chalk.dim("ctrl+c") +
|
||||
chalk.gray(" to clear") +
|
||||
"\n" +
|
||||
chalk.dim("ctrl+c twice") +
|
||||
chalk.gray(" to exit") +
|
||||
"\n" +
|
||||
chalk.dim("ctrl+k") +
|
||||
chalk.gray(" to delete line") +
|
||||
"\n" +
|
||||
chalk.dim("/") +
|
||||
chalk.gray(" for commands") +
|
||||
"\n" +
|
||||
chalk.dim("drop files") +
|
||||
chalk.gray(" to attach");
|
||||
const header = new Text(logo + "\n" + instructions);
|
||||
|
||||
// Setup UI layout
|
||||
this.ui.addChild(header);
|
||||
this.ui.addChild(this.chatContainer);
|
||||
this.ui.addChild(this.statusContainer);
|
||||
this.ui.addChild(this.editorContainer); // Use container that can hold editor or selector
|
||||
this.ui.addChild(this.footer);
|
||||
this.ui.setFocus(this.editor);
|
||||
|
||||
// Set up custom key handlers on the editor
|
||||
this.editor.onEscape = () => {
|
||||
// Intercept Escape key when processing
|
||||
if (this.loadingAnimation && this.onInterruptCallback) {
|
||||
this.onInterruptCallback();
|
||||
}
|
||||
};
|
||||
|
||||
this.editor.onCtrlC = () => {
|
||||
this.handleCtrlC();
|
||||
};
|
||||
|
||||
// Handle editor submission
|
||||
this.editor.onSubmit = (text: string) => {
|
||||
text = text.trim();
|
||||
if (!text) return;
|
||||
|
||||
// Check for /thinking command
|
||||
if (text === "/thinking") {
|
||||
// Show thinking level selector
|
||||
this.showThinkingSelector();
|
||||
this.editor.setText("");
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for /thinking with argument (direct set)
|
||||
if (text.startsWith("/thinking ")) {
|
||||
const level = text.slice("/thinking ".length).trim() as ThinkingLevel;
|
||||
const validLevels: ThinkingLevel[] = ["off", "minimal", "low", "medium", "high"];
|
||||
if (validLevels.includes(level)) {
|
||||
this.agent.setThinkingLevel(level);
|
||||
// Show confirmation message with padding
|
||||
this.chatContainer.addChild(new Text("", 0, 0)); // Blank line before
|
||||
const confirmText = new Text(chalk.blue(`Thinking level set to: ${level}`), 0, 0);
|
||||
this.chatContainer.addChild(confirmText);
|
||||
this.chatContainer.addChild(new Text("", 0, 0)); // Blank line after
|
||||
this.ui.requestRender();
|
||||
this.editor.setText("");
|
||||
return;
|
||||
} else {
|
||||
// Show error message
|
||||
const errorText = new Text(
|
||||
chalk.red(`Invalid thinking level: ${level}. Use: off, minimal, low, medium, high`),
|
||||
1,
|
||||
0,
|
||||
);
|
||||
this.chatContainer.addChild(errorText);
|
||||
this.ui.requestRender();
|
||||
this.editor.setText("");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (this.onInputCallback) {
|
||||
this.onInputCallback(text);
|
||||
}
|
||||
};
|
||||
|
||||
// Start the UI
|
||||
this.ui.start();
|
||||
this.isInitialized = true;
|
||||
}
|
||||
|
||||
async handleEvent(event: import("@mariozechner/pi-agent").AgentEvent, state: AgentState): Promise<void> {
|
||||
if (!this.isInitialized) {
|
||||
await this.init();
|
||||
}
|
||||
|
||||
// Update footer with current stats
|
||||
this.footer.updateState(state);
|
||||
|
||||
switch (event.type) {
|
||||
case "agent_start":
|
||||
// Show loading animation
|
||||
this.editor.disableSubmit = true;
|
||||
// Stop old loader before clearing
|
||||
if (this.loadingAnimation) {
|
||||
this.loadingAnimation.stop();
|
||||
}
|
||||
this.statusContainer.clear();
|
||||
this.loadingAnimation = new Loader(this.ui, "Working... (esc to interrupt)");
|
||||
this.statusContainer.addChild(this.loadingAnimation);
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
|
||||
case "message_start":
|
||||
if (event.message.role === "user") {
|
||||
// Show user message immediately and clear editor
|
||||
this.addMessageToChat(event.message);
|
||||
this.editor.setText("");
|
||||
this.ui.requestRender();
|
||||
} else if (event.message.role === "assistant") {
|
||||
// Create streaming component for assistant messages
|
||||
this.streamingComponent = new StreamingMessageComponent();
|
||||
this.chatContainer.addChild(this.streamingComponent);
|
||||
this.streamingComponent.updateContent(event.message);
|
||||
this.ui.requestRender();
|
||||
}
|
||||
break;
|
||||
|
||||
case "message_update":
|
||||
// Update streaming component
|
||||
if (this.streamingComponent && event.message.role === "assistant") {
|
||||
this.streamingComponent.updateContent(event.message);
|
||||
this.ui.requestRender();
|
||||
}
|
||||
break;
|
||||
|
||||
case "message_end":
|
||||
// Skip user messages (already shown in message_start)
|
||||
if (event.message.role === "user") {
|
||||
break;
|
||||
}
|
||||
if (this.streamingComponent && event.message.role === "assistant") {
|
||||
this.chatContainer.removeChild(this.streamingComponent);
|
||||
this.streamingComponent = null;
|
||||
}
|
||||
// Show final assistant message
|
||||
this.addMessageToChat(event.message);
|
||||
this.ui.requestRender();
|
||||
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();
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool_execution_end": {
|
||||
// Update the existing tool component with the result
|
||||
const component = this.pendingTools.get(event.toolCallId);
|
||||
if (component) {
|
||||
// Update the component with the result
|
||||
component.updateResult({
|
||||
output: typeof event.result === "string" ? event.result : event.result.output,
|
||||
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
|
||||
this.addStatsComponent(this.deferredStats.usage);
|
||||
this.deferredStats = null;
|
||||
}
|
||||
}
|
||||
|
||||
this.ui.requestRender();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case "agent_end":
|
||||
// Stop loading animation
|
||||
if (this.loadingAnimation) {
|
||||
this.loadingAnimation.stop();
|
||||
this.loadingAnimation = null;
|
||||
this.statusContainer.clear();
|
||||
}
|
||||
if (this.streamingComponent) {
|
||||
this.chatContainer.removeChild(this.streamingComponent);
|
||||
this.streamingComponent = null;
|
||||
}
|
||||
this.pendingTools.clear();
|
||||
this.deferredStats = null; // Clear any deferred stats
|
||||
this.editor.disableSubmit = false;
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
private addMessageToChat(message: Message): void {
|
||||
if (message.role === "user") {
|
||||
const userMsg = message as any;
|
||||
// Extract text content from content blocks
|
||||
const textBlocks = userMsg.content.filter((c: any) => c.type === "text");
|
||||
const textContent = textBlocks.map((c: any) => c.text).join("");
|
||||
if (textContent) {
|
||||
// User messages with dark gray background
|
||||
this.chatContainer.addChild(new Markdown(textContent, undefined, undefined, { r: 52, g: 53, b: 65 }));
|
||||
}
|
||||
} else if (message.role === "assistant") {
|
||||
const assistantMsg = message as AssistantMessage;
|
||||
|
||||
// Render content in order
|
||||
for (const content of assistantMsg.content) {
|
||||
if (content.type === "text" && content.text.trim()) {
|
||||
// Assistant text messages with no background - trim the text
|
||||
this.chatContainer.addChild(new Markdown(content.text.trim()));
|
||||
} else if (content.type === "thinking" && content.thinking.trim()) {
|
||||
// Thinking traces in dark gray italic
|
||||
const thinkingText = content.thinking
|
||||
.split("\n")
|
||||
.map((line) => chalk.gray.italic(line))
|
||||
.join("\n");
|
||||
this.chatContainer.addChild(new Text(thinkingText, 1, 0));
|
||||
}
|
||||
}
|
||||
|
||||
// Check if aborted - show after partial content
|
||||
if (assistantMsg.stopReason === "aborted") {
|
||||
// Show red "Aborted" message after partial content
|
||||
const abortedText = new Text(chalk.red("Aborted"));
|
||||
this.chatContainer.addChild(abortedText);
|
||||
return;
|
||||
}
|
||||
|
||||
if (assistantMsg.stopReason === "error") {
|
||||
// Show red error message after partial content
|
||||
const errorMsg = assistantMsg.errorMessage || "Unknown error";
|
||||
const errorText = new Text(chalk.red(`Error: ${errorMsg}`));
|
||||
this.chatContainer.addChild(errorText);
|
||||
return;
|
||||
}
|
||||
|
||||
// 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<string>();
|
||||
for (const content of assistantMsg.content) {
|
||||
if (content.type === "toolCall") {
|
||||
toolCallIds.add(content.id);
|
||||
}
|
||||
}
|
||||
this.deferredStats = { usage: assistantMsg.usage, toolCallIds };
|
||||
} else {
|
||||
// No tool calls - show stats immediately
|
||||
this.addStatsComponent(assistantMsg.usage);
|
||||
}
|
||||
}
|
||||
// Note: tool calls and results are now handled via tool_execution_start/end events
|
||||
}
|
||||
|
||||
private addStatsComponent(usage: any): void {
|
||||
if (!usage) return;
|
||||
|
||||
// Format token counts (similar to web-ui)
|
||||
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)}`);
|
||||
|
||||
if (statsParts.length > 0) {
|
||||
const statsText = new Text(chalk.gray(statsParts.join(" ")), 1, 0);
|
||||
this.chatContainer.addChild(statsText);
|
||||
// Add empty line after stats
|
||||
this.chatContainer.addChild(new Text("", 1, 0));
|
||||
}
|
||||
}
|
||||
|
||||
renderInitialMessages(state: AgentState): void {
|
||||
// Render all existing messages (for --continue mode)
|
||||
// Track assistant messages with their tool calls to show stats after tools
|
||||
const assistantWithTools = new Map<
|
||||
number,
|
||||
{ usage: any; toolCallIds: Set<string>; remainingToolCallIds: Set<string> }
|
||||
>();
|
||||
|
||||
// First pass: identify assistant messages with tool calls
|
||||
for (let i = 0; i < state.messages.length; i++) {
|
||||
const message = state.messages[i];
|
||||
if (message.role === "assistant") {
|
||||
const assistantMsg = message as AssistantMessage;
|
||||
const toolCallIds = new Set<string>();
|
||||
for (const content of assistantMsg.content) {
|
||||
if (content.type === "toolCall") {
|
||||
toolCallIds.add(content.id);
|
||||
}
|
||||
}
|
||||
if (toolCallIds.size > 0) {
|
||||
assistantWithTools.set(i, {
|
||||
usage: assistantMsg.usage,
|
||||
toolCallIds,
|
||||
remainingToolCallIds: new Set(toolCallIds),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Second pass: render messages
|
||||
for (let i = 0; i < state.messages.length; i++) {
|
||||
const message = state.messages[i];
|
||||
|
||||
if (message.role === "user" || message.role === "assistant") {
|
||||
// Temporarily disable deferred stats for initial render
|
||||
const savedDeferredStats = this.deferredStats;
|
||||
this.deferredStats = null;
|
||||
this.addMessageToChat(message);
|
||||
this.deferredStats = savedDeferredStats;
|
||||
} else if (message.role === "toolResult") {
|
||||
// Render tool calls that have already completed
|
||||
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
|
||||
this.addStatsComponent(assistantData.usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
this.ui.requestRender();
|
||||
}
|
||||
|
||||
async getUserInput(): Promise<string> {
|
||||
return new Promise((resolve) => {
|
||||
this.onInputCallback = (text: string) => {
|
||||
this.onInputCallback = undefined;
|
||||
resolve(text);
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
setInterruptCallback(callback: () => void): void {
|
||||
this.onInterruptCallback = callback;
|
||||
}
|
||||
|
||||
private handleCtrlC(): void {
|
||||
// Handle Ctrl+C double-press logic
|
||||
const now = Date.now();
|
||||
const timeSinceLastCtrlC = now - this.lastSigintTime;
|
||||
|
||||
if (timeSinceLastCtrlC < 500) {
|
||||
// Second Ctrl+C within 500ms - exit
|
||||
this.stop();
|
||||
process.exit(0);
|
||||
} else {
|
||||
// First Ctrl+C - clear the editor
|
||||
this.clearEditor();
|
||||
this.lastSigintTime = now;
|
||||
}
|
||||
}
|
||||
|
||||
clearEditor(): void {
|
||||
this.editor.setText("");
|
||||
this.statusContainer.clear();
|
||||
this.ui.requestRender();
|
||||
}
|
||||
|
||||
private showThinkingSelector(): void {
|
||||
const thinkingLevels: SelectItem[] = [
|
||||
{ value: "off", label: "off", description: "No reasoning" },
|
||||
{ value: "minimal", label: "minimal", description: "Very brief reasoning (~1k tokens)" },
|
||||
{ value: "low", label: "low", description: "Light reasoning (~2k tokens)" },
|
||||
{ value: "medium", label: "medium", description: "Moderate reasoning (~8k tokens)" },
|
||||
{ value: "high", label: "high", description: "Deep reasoning (~16k tokens)" },
|
||||
];
|
||||
|
||||
// Create container for the selector with borders
|
||||
const selectorContainer = new Container();
|
||||
|
||||
// Add top border
|
||||
const topBorder = new Text(chalk.blue("─".repeat(50)), 0, 0);
|
||||
selectorContainer.addChild(topBorder);
|
||||
|
||||
// Add selector
|
||||
this.thinkingSelector = new SelectList(thinkingLevels, 5);
|
||||
this.thinkingSelector.onSelect = (item) => {
|
||||
// Apply the selected thinking level
|
||||
const level = item.value as ThinkingLevel;
|
||||
this.agent.setThinkingLevel(level);
|
||||
|
||||
// Show confirmation message with padding and blue color
|
||||
this.chatContainer.addChild(new Text("", 0, 0)); // Blank line before
|
||||
const confirmText = new Text(chalk.blue(`Thinking level set to: ${level}`), 0, 0);
|
||||
this.chatContainer.addChild(confirmText);
|
||||
this.chatContainer.addChild(new Text("", 0, 0)); // Blank line after
|
||||
|
||||
// Hide selector and show editor again
|
||||
this.hideThinkingSelector();
|
||||
this.ui.requestRender();
|
||||
};
|
||||
|
||||
this.thinkingSelector.onCancel = () => {
|
||||
// Just hide the selector
|
||||
this.hideThinkingSelector();
|
||||
this.ui.requestRender();
|
||||
};
|
||||
|
||||
selectorContainer.addChild(this.thinkingSelector);
|
||||
|
||||
// Add bottom border
|
||||
const bottomBorder = new Text(chalk.blue("─".repeat(50)), 0, 0);
|
||||
selectorContainer.addChild(bottomBorder);
|
||||
|
||||
// Replace editor with selector container
|
||||
this.editorContainer.clear();
|
||||
this.editorContainer.addChild(selectorContainer);
|
||||
this.ui.setFocus(this.thinkingSelector);
|
||||
this.ui.requestRender();
|
||||
}
|
||||
|
||||
private hideThinkingSelector(): void {
|
||||
// Replace selector with editor in the container
|
||||
this.editorContainer.clear();
|
||||
this.editorContainer.addChild(this.editor);
|
||||
this.thinkingSelector = null;
|
||||
this.ui.setFocus(this.editor);
|
||||
}
|
||||
|
||||
stop(): void {
|
||||
if (this.loadingAnimation) {
|
||||
this.loadingAnimation.stop();
|
||||
this.loadingAnimation = null;
|
||||
}
|
||||
if (this.isInitialized) {
|
||||
this.ui.stop();
|
||||
this.isInitialized = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue