Smooth UI (#104)

* draft: rewrite single message when streaming (not full terminal)

* having the spinner align better with dots in compact mode

* fix user messages

* handle usage display

* fix formatting

* bash highlighting

---------

Co-authored-by: Nate Woods <[email protected]>

Author: ezynda3

internal/ui/cli.go

@@ -28,6 +28,8 @@ type CLI struct {
 	height           int
 	compactMode      bool   // Add compact mode flag
 	modelName        string // Store current model name
+	lastStreamHeight int    // track how far back we need to move the cursor to overwrite streaming messages
+	usageDisplayed   bool   // track if usage info was displayed after last assistant message
 }
 
 // NewCLI creates a new CLI instance with message container
@@ -64,21 +66,10 @@ func (c *CLI) GetPrompt() (string, error) {
 	// Usage info is now displayed immediately after responses via DisplayUsageAfterResponse()
 	// No need to display it here to avoid duplication
 
-	// Create an enhanced divider with gradient effect
-	theme := GetTheme()
-	dividerStyle := lipgloss.NewStyle().
-		Width(c.width).
-		BorderTop(true).
-		BorderStyle(lipgloss.Border{
-			Top: "━",
-		}).
-		BorderForeground(theme.Border).
-		MarginTop(1).
-		MarginBottom(1).
-		PaddingLeft(2)
-
-	// Render the enhanced input section
-	fmt.Print(dividerStyle.Render(""))
+	c.messageContainer.messages = nil // clear previous messages (they should have been printed already)
+	c.lastStreamHeight = 0            // Reset last stream height for new prompt
+
+	// No divider needed - removed for cleaner appearance
 
 	var prompt string
 	err := huh.NewForm(huh.NewGroup(huh.NewText().
@@ -143,6 +134,10 @@ func (c *CLI) DisplayAssistantMessageWithModel(message, modelName string) error
 
 // DisplayToolCallMessage displays a tool call in progress
 func (c *CLI) DisplayToolCallMessage(toolName, toolArgs string) {
+
+	c.messageContainer.messages = nil // clear previous messages (they should have been printed already)
+	c.lastStreamHeight = 0            // Reset last stream height for new prompt
+
 	var msg UIMessage
 	if c.compactMode {
 		msg = c.compactRenderer.RenderToolCallMessage(toolName, toolArgs, time.Now())
@@ -178,6 +173,8 @@ func (c *CLI) StartStreamingMessage(modelName string) {
 	} else {
 		msg = c.messageRenderer.RenderAssistantMessage("", time.Now(), modelName)
 	}
+	msg.Streaming = true
+	c.lastStreamHeight = 0 // Reset last stream height for new message
 	c.messageContainer.AddMessage(msg)
 	c.displayContainer()
 }
@@ -381,16 +378,49 @@ func (c *CLI) ClearMessages() {
 
 // displayContainer renders and displays the message container
 func (c *CLI) displayContainer() {
-	// Clear screen and display messages
-	fmt.Print("\033[2J\033[H") // Clear screen and move cursor to top
 
 	// Add left padding to the entire container
 	content := c.messageContainer.Render()
+
+	// Check if we're displaying a user message
+	// User messages should not have additional left padding since they're right-aligned
+	// This only applies in non-compact mode
+	paddingLeft := 2
+	if !c.compactMode && len(c.messageContainer.messages) > 0 {
+		lastMessage := c.messageContainer.messages[len(c.messageContainer.messages)-1]
+		if lastMessage.Type == UserMessage {
+			paddingLeft = 0
+		}
+	}
+
 	paddedContent := lipgloss.NewStyle().
-		PaddingLeft(2).
+		PaddingLeft(paddingLeft).
+		Width(c.width). // overwrite (no content) while agent is streaming
 		Render(content)
 
-	fmt.Print(paddedContent)
+	if c.lastStreamHeight > 0 {
+		// Move cursor up by the height of the last streamed message
+		fmt.Printf("\033[%dF", c.lastStreamHeight)
+	} else if c.usageDisplayed {
+		// If we're not overwriting a streaming message but usage was displayed,
+		// move up to account for the usage info (2 lines: content + padding)
+		fmt.Printf("\033[2F")
+		c.usageDisplayed = false
+	}
+
+	fmt.Println(paddedContent)
+
+	// clear message history except the "in-progress" message
+	if len(c.messageContainer.messages) > 0 {
+		// keep the last message, clear the rest (in case of streaming)
+		last := c.messageContainer.messages[len(c.messageContainer.messages)-1]
+		c.messageContainer.messages = []UIMessage{}
+		if last.Streaming {
+			// If the last message is still streaming, we keep it
+			c.messageContainer.messages = append(c.messageContainer.messages, last)
+			c.lastStreamHeight = lipgloss.Height(paddedContent)
+		}
+	}
 }
 
 // UpdateUsage updates the usage tracker with token counts and costs
@@ -487,6 +517,7 @@ func (c *CLI) DisplayUsageAfterResponse() {
 			PaddingTop(1).
 			Render(usageInfo)
 		fmt.Print(paddedUsage)
+		c.usageDisplayed = true
 	}
 }
 

internal/ui/compact_renderer.go

@@ -353,6 +353,11 @@ func (r *CompactRenderer) formatToolResult(result string) string {
 		return ""
 	}
 
+	// Check if this is bash output with stdout/stderr tags
+	if strings.Contains(result, "<stdout>") || strings.Contains(result, "<stderr>") {
+		result = r.formatBashOutput(result)
+	}
+
 	// Calculate available width more conservatively
 	availableWidth := r.width - 28
 	if availableWidth < 40 {
@@ -377,6 +382,74 @@ func (r *CompactRenderer) formatToolResult(result string) string {
 	return strings.Join(lines, "\n")
 }
 
+// formatBashOutput formats bash command output by removing stdout/stderr tags and styling appropriately
+func (r *CompactRenderer) formatBashOutput(result string) string {
+	theme := getTheme()
+
+	// Replace tag pairs with styled content
+	var formattedResult strings.Builder
+	remaining := result
+
+	for {
+		// Find stderr tags
+		stderrStart := strings.Index(remaining, "<stderr>")
+		stderrEnd := strings.Index(remaining, "</stderr>")
+
+		// Find stdout tags
+		stdoutStart := strings.Index(remaining, "<stdout>")
+		stdoutEnd := strings.Index(remaining, "</stdout>")
+
+		// Process whichever comes first
+		if stderrStart != -1 && stderrEnd != -1 && stderrEnd > stderrStart &&
+			(stdoutStart == -1 || stderrStart < stdoutStart) {
+			// Process stderr
+			// Add content before the tag
+			if stderrStart > 0 {
+				formattedResult.WriteString(remaining[:stderrStart])
+			}
+
+			// Extract and style stderr content
+			stderrContent := remaining[stderrStart+8 : stderrEnd]
+			// Trim leading/trailing newlines but preserve internal ones
+			stderrContent = strings.Trim(stderrContent, "\n")
+			if len(stderrContent) > 0 {
+				// Style stderr content with error color, same as non-compact mode
+				styledContent := lipgloss.NewStyle().Foreground(theme.Error).Render(stderrContent)
+				formattedResult.WriteString(styledContent)
+			}
+
+			// Continue with remaining content
+			remaining = remaining[stderrEnd+9:] // Skip past </stderr>
+
+		} else if stdoutStart != -1 && stdoutEnd != -1 && stdoutEnd > stdoutStart {
+			// Process stdout
+			// Add content before the tag
+			if stdoutStart > 0 {
+				formattedResult.WriteString(remaining[:stdoutStart])
+			}
+
+			// Extract stdout content (no special styling needed)
+			stdoutContent := remaining[stdoutStart+8 : stdoutEnd]
+			// Trim leading/trailing newlines but preserve internal ones
+			stdoutContent = strings.Trim(stdoutContent, "\n")
+			if len(stdoutContent) > 0 {
+				formattedResult.WriteString(stdoutContent)
+			}
+
+			// Continue with remaining content
+			remaining = remaining[stdoutEnd+9:] // Skip past </stdout>
+
+		} else {
+			// No more tags, add remaining content
+			formattedResult.WriteString(remaining)
+			break
+		}
+	}
+
+	// Trim any leading/trailing whitespace from the final result
+	return strings.TrimSpace(formattedResult.String())
+}
+
 // determineResultType determines the display type for tool results
 func (r *CompactRenderer) determineResultType(toolName, result string) string {
 	toolName = strings.ToLower(toolName)
@@ -386,7 +459,7 @@ func (r *CompactRenderer) determineResultType(toolName, result string) string {
 		return "Text"
 	case strings.Contains(toolName, "write"):
 		return "Write"
-	case strings.Contains(toolName, "bash") || strings.Contains(toolName, "command"):
+	case strings.Contains(toolName, "bash") || strings.Contains(toolName, "command") || strings.Contains(toolName, "shell") || toolName == "run_shell_cmd":
 		return "Bash"
 	case strings.Contains(toolName, "list") || strings.Contains(toolName, "ls"):
 		return "List"

internal/ui/factory.go

@@ -75,6 +75,8 @@ func SetupCLI(opts *CLISetupOptions) (*CLI, error) {
 		}
 	}
 
+	fmt.Println("")
+
 	// Display model info
 	if provider != "unknown" && model != "unknown" {
 		cli.DisplayInfo(fmt.Sprintf("Model loaded: %s (%s)", provider, model))
@@ -89,5 +91,10 @@ func SetupCLI(opts *CLISetupOptions) (*CLI, error) {
 	tools := opts.Agent.GetTools()
 	cli.DisplayInfo(fmt.Sprintf("Loaded %d tools from MCP servers", len(tools)))
 
+	// Display usage information (for both streaming and non-streaming)
+	if !opts.Quiet && cli != nil {
+		cli.DisplayUsageAfterResponse()
+	}
+
 	return cli, nil
 }

internal/ui/messages.go

@@ -30,6 +30,7 @@ type UIMessage struct {
 	Height    int
 	Content   string
 	Timestamp time.Time
+	Streaming bool
 }
 
 // Helper functions to get theme colors
@@ -411,7 +412,7 @@ func (r *MessageRenderer) formatToolResult(toolName, result string, width int) s
 	}
 
 	// Format bash/command output with better formatting
-	if strings.Contains(toolName, "bash") || strings.Contains(toolName, "command") {
+	if strings.Contains(toolName, "bash") || strings.Contains(toolName, "command") || strings.Contains(toolName, "shell") || toolName == "run_shell_cmd" {
 		theme := getTheme()
 
 		// Split result into sections if it contains both stdout and stderr
@@ -438,49 +439,71 @@ func (r *MessageRenderer) formatToolResult(toolName, result string, width int) s
 func (r *MessageRenderer) formatBashOutput(result string, width int, theme Theme) string {
 	baseStyle := lipgloss.NewStyle()
 
-	// Parse the output sections
-	lines := strings.Split(result, "\n")
-	var formattedLines []string
-	currentSection := ""
-
-	for _, line := range lines {
-		if strings.HasPrefix(line, "<stdout>") {
-			currentSection = "stdout"
-			if len(strings.TrimSpace(strings.TrimPrefix(line, "<stdout>"))) > 0 {
-				// If there's content on the same line as <stdout>
-				content := strings.TrimSpace(strings.TrimPrefix(line, "<stdout>"))
-				formattedLines = append(formattedLines, content)
+	// Replace tag pairs with styled content
+	var formattedResult strings.Builder
+	remaining := result
+
+	for {
+		// Find stderr tags
+		stderrStart := strings.Index(remaining, "<stderr>")
+		stderrEnd := strings.Index(remaining, "</stderr>")
+
+		// Find stdout tags
+		stdoutStart := strings.Index(remaining, "<stdout>")
+		stdoutEnd := strings.Index(remaining, "</stdout>")
+
+		// Process whichever comes first
+		if stderrStart != -1 && stderrEnd != -1 && stderrEnd > stderrStart &&
+			(stdoutStart == -1 || stderrStart < stdoutStart) {
+			// Process stderr
+			// Add content before the tag
+			if stderrStart > 0 {
+				formattedResult.WriteString(remaining[:stderrStart])
 			}
-			continue
-		} else if strings.HasPrefix(line, "<stderr>") {
-			currentSection = "stderr"
-			if len(strings.TrimSpace(strings.TrimPrefix(line, "<stderr>"))) > 0 {
-				// If there's content on the same line as <stderr>
-				content := strings.TrimSpace(strings.TrimPrefix(line, "<stderr>"))
-				styledLine := baseStyle.Foreground(theme.Error).Render(content)
-				formattedLines = append(formattedLines, styledLine)
+			// Extract and style stderr content
+			stderrContent := remaining[stderrStart+8 : stderrEnd]
+			// Trim leading/trailing newlines but preserve internal ones
+			stderrContent = strings.Trim(stderrContent, "\n")
+			if len(stderrContent) > 0 {
+				styledContent := baseStyle.Foreground(theme.Error).Render(stderrContent)
+				formattedResult.WriteString(styledContent)
 			}
-			continue
-		} else if line == "" {
-			// Preserve empty lines but don't add extra spacing
-			formattedLines = append(formattedLines, "")
-			continue
-		}
 
-		// Regular content line
-		if currentSection == "stderr" {
-			styledLine := baseStyle.Foreground(theme.Error).Render(line)
-			formattedLines = append(formattedLines, styledLine)
+			// Continue with remaining content
+			remaining = remaining[stderrEnd+9:] // Skip past </stderr>
+
+		} else if stdoutStart != -1 && stdoutEnd != -1 && stdoutEnd > stdoutStart {
+			// Process stdout
+			// Add content before the tag
+			if stdoutStart > 0 {
+				formattedResult.WriteString(remaining[:stdoutStart])
+			}
+
+			// Extract stdout content (no special styling needed)
+			stdoutContent := remaining[stdoutStart+8 : stdoutEnd]
+			// Trim leading/trailing newlines but preserve internal ones
+			stdoutContent = strings.Trim(stdoutContent, "\n")
+			if len(stdoutContent) > 0 {
+				formattedResult.WriteString(stdoutContent)
+			}
+
+			// Continue with remaining content
+			remaining = remaining[stdoutEnd+9:] // Skip past </stdout>
+
 		} else {
-			// stdout or no section - use normal muted color
-			formattedLines = append(formattedLines, line)
+			// No more tags, add remaining content
+			formattedResult.WriteString(remaining)
+			break
 		}
 	}
 
+	// Trim any leading/trailing whitespace from the final result
+	finalResult := strings.TrimSpace(formattedResult.String())
+
 	return baseStyle.
 		Width(width).
 		Foreground(theme.Muted).
-		Render(strings.Join(formattedLines, "\n"))
+		Render(finalResult)
 }
 
 // truncateText truncates text to fit within the specified width
@@ -563,6 +586,7 @@ func (c *MessageContainer) UpdateLastMessage(content string) {
 			renderer := NewMessageRenderer(c.width, false)
 			newMsg = renderer.RenderAssistantMessage(content, lastMsg.Timestamp, c.modelName)
 		}
+		newMsg.Streaming = lastMsg.Streaming // Preserve streaming state
 		c.messages[lastIdx] = newMsg
 	}
 }
@@ -608,12 +632,14 @@ func (c *MessageContainer) Render() string {
 		}
 	}
 
-	return lipgloss.NewStyle().
-		Width(c.width).
-		PaddingBottom(1).
-		Render(
-			lipgloss.JoinVertical(lipgloss.Top, parts...),
-		)
+	style := lipgloss.NewStyle().
+		Width(c.width)
+
+	// No padding needed between messages
+
+	return style.Render(
+		lipgloss.JoinVertical(lipgloss.Top, parts...),
+	)
 }
 
 // renderEmptyState renders an enhanced initial empty state
@@ -695,7 +721,7 @@ func (c *MessageContainer) renderCompactMessages() string {
 		lines = append(lines, msg.Content)
 	}
 
-	return strings.Join(lines, "\n") + "\n"
+	return strings.Join(lines, "\n")
 }
 
 // renderCompactEmptyState renders a simple empty state for compact mode