Merge pull request docker#1496 from krissetto/improve-diff-rendering

Cache file reads for diff rendering

Author: dgageot

pkg/tui/components/reasoningblock/reasoningblock.go

@@ -519,7 +519,13 @@ func (m *Model) renderCollapsed() string {
 	if len(visibleTools) > 0 {
 		parts = append(parts, "") // blank line before tools
 		for _, entry := range visibleTools {
-			toolView := entry.view.View()
+			// Prefer CollapsedView() for simplified rendering in collapsed state
+			var toolView string
+			if cv, ok := entry.view.(layout.CollapsedViewer); ok {
+				toolView = cv.CollapsedView()
+			} else {
+				toolView = entry.view.View()
+			}
 			if entry.fadeProgress > 0 {
 				// Strip existing ANSI codes and apply faded color based on progress
 				// (wrapping styled content doesn't override inner colors)

pkg/tui/components/tool/editfile/editfile.go

@@ -17,7 +17,7 @@ type ToggleDiffViewMsg struct{}
 
 // New creates the edit_file tool UI model.
 func New(msg *types.Message, sessionState *service.SessionState) layout.Model {
-	return toolcommon.NewBase(msg, sessionState, render)
+	return toolcommon.NewBaseWithCollapsed(msg, sessionState, render, renderCollapsed)
 }
 
 // render displays the edit_file tool output in the TUI.
@@ -101,3 +101,51 @@ func render(
 
 	return content
 }
+
+// renderCollapsed renders a simplified view for collapsed reasoning blocks.
+// Shows only the file path and +N / -M line counts.
+func renderCollapsed(
+	msg *types.Message,
+	s spinner.Spinner,
+	_ *service.SessionState,
+	width,
+	_ int,
+) string {
+	var args builtin.EditFileArgs
+	if err := json.Unmarshal([]byte(msg.ToolCall.Function.Arguments), &args); err != nil {
+		return ""
+	}
+
+	// Error state
+	if msg.ToolStatus == types.ToolStatusError {
+		if msg.Content == "" {
+			return ""
+		}
+		line := fmt.Sprintf(
+			"%s%s %s",
+			toolcommon.Icon(msg, s),
+			styles.ToolNameError.Render(msg.ToolDefinition.DisplayName()),
+			styles.ToolErrorMessageStyle.Render(msg.Content),
+		)
+		return styles.BaseStyle.MaxWidth(width).Render(line)
+	}
+
+	// Count added/removed lines
+	added, removed := countDiffLines(msg.ToolCall, msg.ToolStatus)
+	var diffSummary string
+	if added > 0 || removed > 0 {
+		addStr := styles.DiffAddStyle.Render(fmt.Sprintf("+%d", added))
+		remStr := styles.DiffRemoveStyle.Render(fmt.Sprintf("-%d", removed))
+		diffSummary = fmt.Sprintf(" %s / %s", addStr, remStr)
+	}
+
+	line := fmt.Sprintf(
+		"%s%s %s%s",
+		toolcommon.Icon(msg, s),
+		styles.ToolName.Render(msg.ToolDefinition.DisplayName()),
+		styles.ToolMessageStyle.Render(toolcommon.ShortenPath(args.Path)),
+		diffSummary,
+	)
+
+	return styles.BaseStyle.MaxWidth(width).Render(line)
+}

pkg/tui/components/tool/editfile/render.go

@@ -27,7 +27,24 @@ const (
 	minWidth     = 80
 )
 
+type toolRenderCache struct {
+	// Line counts - computed once, never change
+	added       int
+	removed     int
+	lineCounted bool
+
+	// Rendered output - invalidated when width/splitView/status changes
+	rendered       string
+	renderCached   bool
+	renderedWidth  int
+	renderedSplit  bool
+	renderedStatus types.ToolStatus
+}
+
 var (
+	cache   = make(map[string]*toolRenderCache) // keyed by toolCallID
+	cacheMu sync.RWMutex
+
 	lexerCache   = make(map[string]chroma.Lexer)
 	lexerCacheMu sync.RWMutex
 )
@@ -44,7 +61,53 @@ type linePair struct {
 	newLineNum int
 }
 
+func getOrCreateCache(toolCallID string) *toolRenderCache {
+	cacheMu.RLock()
+	if c, ok := cache[toolCallID]; ok {
+		cacheMu.RUnlock()
+		return c
+	}
+	cacheMu.RUnlock()
+
+	cacheMu.Lock()
+	defer cacheMu.Unlock()
+	// Double-check after acquiring write lock
+	if c, ok := cache[toolCallID]; ok {
+		return c
+	}
+	c := &toolRenderCache{}
+	cache[toolCallID] = c
+	return c
+}
+
 func renderEditFile(toolCall tools.ToolCall, width int, splitView bool, toolStatus types.ToolStatus) string {
+	c := getOrCreateCache(toolCall.ID)
+
+	cacheMu.RLock()
+	if c.renderCached &&
+		c.renderedWidth == width &&
+		c.renderedSplit == splitView &&
+		c.renderedStatus == toolStatus {
+		result := c.rendered
+		cacheMu.RUnlock()
+		return result
+	}
+	cacheMu.RUnlock()
+
+	result := renderEditFileUncached(toolCall, width, splitView, toolStatus)
+
+	cacheMu.Lock()
+	c.rendered = result
+	c.renderCached = true
+	c.renderedWidth = width
+	c.renderedSplit = splitView
+	c.renderedStatus = toolStatus
+	cacheMu.Unlock()
+
+	return result
+}
+
+func renderEditFileUncached(toolCall tools.ToolCall, width int, splitView bool, toolStatus types.ToolStatus) string {
 	var args builtin.EditFileArgs
 	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
 		return ""
@@ -71,6 +134,56 @@ func renderEditFile(toolCall tools.ToolCall, width int, splitView bool, toolStat
 	return output.String()
 }
 
+// countDiffLines returns the number of added and removed lines for the edit.
+// Results are cached per tool call since arguments are immutable.
+func countDiffLines(toolCall tools.ToolCall, _ types.ToolStatus) (added, removed int) {
+	c := getOrCreateCache(toolCall.ID)
+
+	cacheMu.RLock()
+	if c.lineCounted {
+		added, removed = c.added, c.removed
+		cacheMu.RUnlock()
+		return added, removed
+	}
+	cacheMu.RUnlock()
+
+	added, removed = countDiffLinesUncached(toolCall)
+
+	cacheMu.Lock()
+	c.added = added
+	c.removed = removed
+	c.lineCounted = true
+	cacheMu.Unlock()
+
+	return added, removed
+}
+
+func countDiffLinesUncached(toolCall tools.ToolCall) (added, removed int) {
+	var args builtin.EditFileArgs
+	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
+		return 0, 0
+	}
+
+	for _, edit := range args.Edits {
+		edits := udiff.Strings(edit.OldText, edit.NewText)
+		diff, err := udiff.ToUnifiedDiff("old", "new", edit.OldText, edits, 0)
+		if err != nil {
+			continue
+		}
+		for _, hunk := range diff.Hunks {
+			for _, line := range hunk.Lines {
+				switch line.Kind {
+				case udiff.Insert:
+					added++
+				case udiff.Delete:
+					removed++
+				}
+			}
+		}
+	}
+	return added, removed
+}
+
 func computeDiff(path, oldText, newText string, toolStatus types.ToolStatus) []*udiff.Hunk {
 	currentContent, err := os.ReadFile(path)
 	if err != nil {
@@ -91,7 +204,6 @@ func computeDiff(path, oldText, newText string, toolStatus types.ToolStatus) []*
 		oldContent = strings.Replace(newContent, newText, oldText, 1)
 	}
 
-	// Now compute diff between old and new
 	edits := udiff.Strings(oldContent, newContent)
 
 	diff, err := udiff.ToUnifiedDiff("old", "new", oldContent, edits, 3)

pkg/tui/components/toolcommon/base.go

@@ -14,6 +14,9 @@ import (
 // It receives the message, spinner, session state, and available width/height.
 type Renderer func(msg *types.Message, s spinner.Spinner, sessionState *service.SessionState, width, height int) string
 
+// CollapsedRenderer is a function that renders a simplified view for collapsed reasoning blocks.
+type CollapsedRenderer func(msg *types.Message, s spinner.Spinner, sessionState *service.SessionState, width, height int) string
+
 // Base provides common boilerplate for tool components.
 // It handles spinner management, sizing, and delegates rendering to a custom function.
 type Base struct {
@@ -23,6 +26,7 @@ type Base struct {
 	height            int
 	sessionState      *service.SessionState
 	render            Renderer
+	collapsedRenderer CollapsedRenderer
 	spinnerRegistered bool // tracks whether spinner is registered with coordinator
 }
 
@@ -38,6 +42,19 @@ func NewBase(msg *types.Message, sessionState *service.SessionState, render Rend
 	}
 }
 
+// NewBaseWithCollapsed creates a new base tool component with both regular and collapsed renderers.
+func NewBaseWithCollapsed(msg *types.Message, sessionState *service.SessionState, render Renderer, collapsedRender CollapsedRenderer) *Base {
+	return &Base{
+		message:           msg,
+		spinner:           spinner.New(spinner.ModeSpinnerOnly, styles.SpinnerDotsAccentStyle),
+		width:             80,
+		height:            1,
+		sessionState:      sessionState,
+		render:            render,
+		collapsedRenderer: collapsedRender,
+	}
+}
+
 // Message returns the tool message.
 func (b *Base) Message() *types.Message {
 	return b.message
@@ -107,6 +124,15 @@ func (b *Base) View() string {
 	return b.render(b.message, b.spinner, b.sessionState, b.width, b.height)
 }
 
+// CollapsedView returns a simplified view for use in collapsed reasoning blocks.
+// Falls back to the regular View() if no collapsed renderer is provided.
+func (b *Base) CollapsedView() string {
+	if b.collapsedRenderer != nil {
+		return b.collapsedRenderer(b.message, b.spinner, b.sessionState, b.width, b.height)
+	}
+	return b.View()
+}
+
 func (b *Base) isSpinnerActive() bool {
 	return b.message.ToolStatus == types.ToolStatusPending ||
 		b.message.ToolStatus == types.ToolStatusRunning

pkg/tui/core/layout/layout.go

@@ -34,3 +34,9 @@ type Model interface {
 	View() string
 	Sizeable
 }
+
+// CollapsedViewer is implemented by components that provide a simplified view
+// for use in collapsed reasoning blocks.
+type CollapsedViewer interface {
+	CollapsedView() string
+}