Merge pull request #1412 from krissetto/collapsible-reasoning

Collapsible reasoning in TUI

Author: dgageot

pkg/runtime/runtime.go

@@ -1033,8 +1033,13 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 	var actualModelEventEmitted bool
 	var messageUsage *chat.Usage
 	modelID := getAgentModelID(a)
-	// Track which tool call indices we've already emitted partial events for
-	emittedPartialEvents := make(map[string]bool)
+
+	toolCallIndex := make(map[string]int)   // toolCallID -> index in toolCalls slice
+	emittedPartial := make(map[string]bool) // toolCallID -> whether we've emitted a partial event
+	toolDefMap := make(map[string]tools.Tool, len(agentTools))
+	for _, t := range agentTools {
+		toolDefMap[t.Name] = t
+	}
 
 	for {
 		response, err := stream.Recv()
@@ -1103,63 +1108,39 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 		// Handle tool calls
 		if len(choice.Delta.ToolCalls) > 0 {
 			// Process each tool call delta
-			for _, deltaToolCall := range choice.Delta.ToolCalls {
-				// Find existing tool call by ID, or create a new one
-				idx := -1
-				for i, toolCall := range toolCalls {
-					if toolCall.ID == deltaToolCall.ID {
-						idx = i
-						break
-					}
-				}
-
-				// If tool call doesn't exist yet, append it
-				if idx == -1 {
+			for _, delta := range choice.Delta.ToolCalls {
+				idx, exists := toolCallIndex[delta.ID]
+				if !exists {
 					idx = len(toolCalls)
+					toolCallIndex[delta.ID] = idx
 					toolCalls = append(toolCalls, tools.ToolCall{
-						ID:   deltaToolCall.ID,
-						Type: deltaToolCall.Type,
+						ID:   delta.ID,
+						Type: delta.Type,
 					})
 				}
 
-				// Check if we should emit a partial event for this tool call
-				// We want to emit when we first get the function name
-				shouldEmitPartial := !emittedPartialEvents[deltaToolCall.ID] &&
-					deltaToolCall.Function.Name != "" &&
-					toolCalls[idx].Function.Name == "" // Don't emit if we already have the name
+				tc := &toolCalls[idx]
 
-				// Update fields based on what's in the delta
-				if deltaToolCall.ID != "" {
-					toolCalls[idx].ID = deltaToolCall.ID
-				}
-				if deltaToolCall.Type != "" {
-					toolCalls[idx].Type = deltaToolCall.Type
+				// Track if we're learning the name for the first time
+				learningName := delta.Function.Name != "" && tc.Function.Name == ""
+
+				// Update fields from delta
+				if delta.Type != "" {
+					tc.Type = delta.Type
 				}
-				if deltaToolCall.Function.Name != "" {
-					toolCalls[idx].Function.Name = deltaToolCall.Function.Name
+				if delta.Function.Name != "" {
+					tc.Function.Name = delta.Function.Name
 				}
-				if deltaToolCall.Function.Arguments != "" {
-					if toolCalls[idx].Function.Arguments == "" {
-						toolCalls[idx].Function.Arguments = deltaToolCall.Function.Arguments
-					} else {
-						toolCalls[idx].Function.Arguments += deltaToolCall.Function.Arguments
-					}
-					// Emit if we get more arguments
-					shouldEmitPartial = true
+				if delta.Function.Arguments != "" {
+					tc.Function.Arguments += delta.Function.Arguments
 				}
 
-				// Emit PartialToolCallEvent when we first get the function name
-				if shouldEmitPartial {
-					// TODO: clean this up, it's gross
-					tool := tools.Tool{}
-					for _, t := range agentTools {
-						if t.Name == toolCalls[idx].Function.Name {
-							tool = t
-							break
-						}
+				// Emit PartialToolCall once we have a name, and on subsequent argument deltas
+				if tc.Function.Name != "" && (learningName || delta.Function.Arguments != "") {
+					if !emittedPartial[delta.ID] || delta.Function.Arguments != "" {
+						events <- PartialToolCall(*tc, toolDefMap[tc.Function.Name], a.Name())
+						emittedPartial[delta.ID] = true
 					}
-					events <- PartialToolCall(toolCalls[idx], tool, a.Name())
-					emittedPartialEvents[deltaToolCall.ID] = true
 				}
 			}
 			continue

pkg/tui/components/message/message.go

@@ -108,40 +108,6 @@ func (mv *messageModel) Render(width int) string {
 		}
 
 		return mv.senderPrefix(msg.Sender) + messageStyle.Render(rendered)
-	case types.MessageTypeAssistantReasoning:
-		if msg.Content == "" {
-			return mv.spinner.View()
-		}
-
-		messageStyle := styles.AssistantMessageStyle
-		thinkingStyle := styles.MutedStyle.Italic(true)
-
-		rendered, err := markdown.NewRenderer(width - messageStyle.GetHorizontalFrameSize()).Render(msg.Content)
-		if err != nil {
-			rendered = msg.Content
-		}
-
-		// Strip ANSI so muted style applies uniformly, and trim trailing whitespace.
-		// Unlike regular content where markdown ANSI output goes directly to messageStyle,
-		// here we strip ANSI which exposes raw trailing newlines from markdown that would
-		// otherwise be handled differently by lipgloss when embedded in ANSI sequences.
-		clean := strings.TrimRight(stripANSI(rendered), "\n\r\t ")
-
-		// Show "Thinking:" badge only when starting a new thinking block after content
-		var text string
-		if mv.continuingThinking(msg) {
-			text = thinkingStyle.Render(clean)
-		} else {
-			text = styles.ThinkingBadgeStyle.Render("Thinking:") + "\n\n" + thinkingStyle.Render(clean)
-		}
-
-		styledContent := messageStyle.Render(text)
-
-		if mv.sameAgentAsPrevious(msg) {
-			return styledContent
-		}
-
-		return mv.senderPrefix(msg.Sender) + styledContent
 	case types.MessageTypeShellOutput:
 		if rendered, err := markdown.NewRenderer(width).Render(fmt.Sprintf("```console\n%s\n```", msg.Content)); err == nil {
 			return rendered
@@ -187,23 +153,7 @@ func (mv *messageModel) sameAgentAsPrevious(msg *types.Message) bool {
 	}
 	switch mv.previous.Type {
 	case types.MessageTypeAssistant,
-		types.MessageTypeAssistantReasoning,
-		types.MessageTypeToolCall,
-		types.MessageTypeToolResult:
-		return true
-	default:
-		return false
-	}
-}
-
-// continuingThinking returns true if we're continuing a thinking flow
-// (previous was thinking or tool call, not content)
-func (mv *messageModel) continuingThinking(msg *types.Message) bool {
-	if mv.previous == nil || mv.previous.Sender != msg.Sender {
-		return false
-	}
-	switch mv.previous.Type {
-	case types.MessageTypeAssistantReasoning,
+		types.MessageTypeAssistantReasoningBlock,
 		types.MessageTypeToolCall,
 		types.MessageTypeToolResult:
 		return true