Throttle and merge incoming events before sending them to the TUI

Sometimes the LLMs just send too many tokens per second, we don't want
to re-render completely for each token in that case. So we throttle and
merge partial tool call, reasoning and agent choice events.

This change brings the CPU usage on fast LLMs down from ~130% CPU to
~25-30%.

Signed-off-by: Djordje Lukic <djordje.lukic@docker.com>

Author: rumpl

pkg/app/app.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"os/exec"
 	"strings"
+	"time"
 
 	tea "github.com/charmbracelet/bubbletea/v2"
 
@@ -13,24 +14,26 @@ import (
 )
 
 type App struct {
-	title         string
-	agentFilename string
-	runtime       runtime.Runtime
-	team          *team.Team
-	session       *session.Session
-	firstMessage  *string
-	events        chan tea.Msg
+	title            string
+	agentFilename    string
+	runtime          runtime.Runtime
+	team             *team.Team
+	session          *session.Session
+	firstMessage     *string
+	events           chan tea.Msg
+	throttleDuration time.Duration
 }
 
 func New(title, agentFilename string, rt runtime.Runtime, agents *team.Team, sess *session.Session, firstMessage *string) *App {
 	return &App{
-		title:         title,
-		agentFilename: agentFilename,
-		runtime:       rt,
-		team:          agents,
-		session:       sess,
-		firstMessage:  firstMessage,
-		events:        make(chan tea.Msg, 128),
+		title:            title,
+		agentFilename:    agentFilename,
+		runtime:          rt,
+		team:             agents,
+		session:          sess,
+		firstMessage:     firstMessage,
+		events:           make(chan tea.Msg, 128),
+		throttleDuration: 50 * time.Millisecond, // Throttle rapid events
 	}
 }
 
@@ -65,11 +68,12 @@ func (a *App) Run(ctx context.Context, message string) {
 }
 
 func (a *App) Subscribe(ctx context.Context, program *tea.Program) {
+	throttledChan := a.throttleEvents(ctx, a.events)
 	for {
 		select {
 		case <-ctx.Done():
 			return
-		case msg, ok := <-a.events:
+		case msg, ok := <-throttledChan:
 			if !ok {
 				return
 			}
@@ -91,3 +95,151 @@ func (a *App) ResumeStartOAuth(confirmation bool) {
 		a.runtime.ResumeStartAuthorizationFlow(context.Background(), confirmation)
 	}
 }
+
+// throttleEvents buffers and merges rapid events to prevent UI flooding
+func (a *App) throttleEvents(ctx context.Context, in <-chan tea.Msg) <-chan tea.Msg {
+	out := make(chan tea.Msg, 128)
+
+	go func() {
+		defer close(out)
+
+		var buffer []tea.Msg
+		ticker := time.NewTicker(a.throttleDuration)
+		defer ticker.Stop()
+
+		flush := func() {
+			if len(buffer) == 0 {
+				return
+			}
+
+			// Merge events if possible
+			merged := a.mergeEvents(buffer)
+			for _, msg := range merged {
+				select {
+				case out <- msg:
+				case <-ctx.Done():
+					return
+				}
+			}
+			buffer = buffer[:0]
+		}
+
+		for {
+			select {
+			case <-ctx.Done():
+				flush()
+				return
+
+			case msg, ok := <-in:
+				if !ok {
+					flush()
+					return
+				}
+
+				// Check if this event type should be throttled
+				if a.shouldThrottle(msg) {
+					buffer = append(buffer, msg)
+				} else {
+					// Pass through immediately for important events
+					flush() // Flush any buffered events first
+					select {
+					case out <- msg:
+					case <-ctx.Done():
+						return
+					}
+				}
+
+			case <-ticker.C:
+				flush()
+			}
+		}
+	}()
+
+	return out
+}
+
+// shouldThrottle determines if an event should be buffered/throttled
+func (a *App) shouldThrottle(msg tea.Msg) bool {
+	switch msg.(type) {
+	case *runtime.AgentChoiceEvent:
+		return true
+	case *runtime.AgentChoiceReasoningEvent:
+		return true
+	case *runtime.PartialToolCallEvent:
+		return true
+	default:
+		return false
+	}
+}
+
+// mergeEvents merges consecutive similar events to reduce UI updates
+func (a *App) mergeEvents(events []tea.Msg) []tea.Msg {
+	if len(events) == 0 {
+		return events
+	}
+
+	var result []tea.Msg
+
+	// Group events by type and merge
+	for i := 0; i < len(events); i++ {
+		current := events[i]
+
+		switch ev := current.(type) {
+		case *runtime.AgentChoiceEvent:
+			// Merge consecutive AgentChoiceEvents with same agent
+			merged := ev
+			for i+1 < len(events) {
+				if next, ok := events[i+1].(*runtime.AgentChoiceEvent); ok && next.AgentName == ev.AgentName {
+					// Concatenate content
+					merged = &runtime.AgentChoiceEvent{
+						Type:         ev.Type,
+						Content:      merged.Content + next.Content,
+						AgentContext: ev.AgentContext,
+					}
+					i++
+				} else {
+					break
+				}
+			}
+			result = append(result, merged)
+
+		case *runtime.AgentChoiceReasoningEvent:
+			// Merge consecutive AgentChoiceReasoningEvents with same agent
+			merged := ev
+			for i+1 < len(events) {
+				if next, ok := events[i+1].(*runtime.AgentChoiceReasoningEvent); ok && next.AgentName == ev.AgentName {
+					// Concatenate content
+					merged = &runtime.AgentChoiceReasoningEvent{
+						Type:         ev.Type,
+						Content:      merged.Content + next.Content,
+						AgentContext: ev.AgentContext,
+					}
+					i++
+				} else {
+					break
+				}
+			}
+			result = append(result, merged)
+
+		case *runtime.PartialToolCallEvent:
+			// For PartialToolCallEvent, keep only the latest one per tool call ID
+			// Check if there's a newer one in the buffer
+			var latest *runtime.PartialToolCallEvent = ev
+			for j := i + 1; j < len(events); j++ {
+				if next, ok := events[j].(*runtime.PartialToolCallEvent); ok {
+					if next.ToolCall.ID == ev.ToolCall.ID {
+						latest = next
+						i = j // Skip to this position
+					}
+				}
+			}
+			result = append(result, latest)
+
+		default:
+			// Pass through other events as-is
+			result = append(result, current)
+		}
+	}
+
+	return result
+}