Handle and print reasoning tokens (#111)

rumpl · web-flow · commit 6b68057eb2a6 · 2025-09-05T14:50:47.000+02:00
The runtime will emit a new event for thinking models `agent_choice_reasoning`. Another change is in the `agent_choice` event, it now only sends the content as a string. <img width="1842" height="1189" alt="Screenshot 2025-09-05 at 14 15 26" src="https://github.com/user-attachments/assets/42306349-7421-474e-96d7-9fa425d45cf5" /> Signed-off-by: Djordje Lukic <djordje.lukic@docker.com>
diff --git a/cmd/root/new.go b/cmd/root/new.go
@@ -107,7 +107,7 @@ func NewNewCmd() *cobra.Command {
 						fmt.Println()
 						llmIsTyping = true
 					}
-					fmt.Printf("%s", e.Choice.Delta.Content)
+					fmt.Printf("%s", e.Content)
 				case *runtime.ToolCallEvent:
 					if llmIsTyping {
 						fmt.Println()
diff --git a/cmd/root/run.go b/cmd/root/run.go
@@ -321,7 +321,7 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt *runtime.Runtim
 					}
 					llmIsTyping = true
 				}
-				fmt.Printf("%s", e.Choice.Delta.Content)
+				fmt.Printf("%s", e.Content)
 			case *runtime.ToolCallConfirmationEvent:
 				if llmIsTyping {
 					fmt.Println()
diff --git a/internal/tui/components/message/message.go b/internal/tui/components/message/message.go
@@ -45,17 +45,15 @@ func New(msg *types.Message, renderer *glamour.TermRenderer) Model {
 
 // Init initializes the message view
 func (mv *messageModel) Init() tea.Cmd {
-	// Start spinner for empty assistant messages
-	if mv.message.Type == types.MessageTypeAssistant && mv.message.Content == "" {
+	if mv.message.Type == types.MessageTypeSpinner {
 		return mv.spinner.Tick
 	}
 	return nil
 }
 
 // Update handles messages and updates the message view state
 func (mv *messageModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
-	// Handle spinner updates for empty assistant messages
-	if mv.message.Type == types.MessageTypeAssistant && mv.message.Content == "" {
+	if mv.message.Type == types.MessageTypeSpinner {
 		var cmd tea.Cmd
 		mv.spinner, cmd = mv.spinner.Update(msg)
 		return mv, cmd
@@ -75,6 +73,8 @@ func (mv *messageModel) View() string {
 func (mv *messageModel) Render(int) string {
 	msg := mv.message
 	switch msg.Type {
+	case types.MessageTypeSpinner:
+		return mv.spinner.View()
 	case types.MessageTypeUser:
 		if rendered, err := mv.renderer.Render("> " + msg.Content); err == nil {
 			return strings.TrimRight(rendered, "\n\r\t ")
@@ -92,6 +92,12 @@ func (mv *messageModel) Render(int) string {
 		}
 
 		return strings.TrimRight(rendered, "\n\r\t ")
+	case types.MessageTypeAssistantReasoning:
+		if msg.Content == "" {
+			return mv.spinner.View()
+		}
+		text := senderPrefix(msg.Sender) + msg.Content
+		return styles.MutedStyle.Italic(true).Render("Thinking: " + text)
 	case types.MessageTypeShellOutput:
 		if rendered, err := mv.renderer.Render(fmt.Sprintf("```console\n%s\n```", msg.Content)); err == nil {
 			return strings.TrimRight(rendered, "\n\r\t ")
diff --git a/internal/tui/components/messages/messages.go b/internal/tui/components/messages/messages.go
@@ -17,6 +17,7 @@ import (
 	"github.com/docker/cagent/internal/tui/core"
 	"github.com/docker/cagent/internal/tui/core/layout"
 	"github.com/docker/cagent/internal/tui/types"
+	"github.com/docker/cagent/pkg/runtime"
 	"github.com/docker/cagent/pkg/tools"
 )
 
@@ -32,8 +33,8 @@ type Model interface {
 	AddAssistantMessage() tea.Cmd
 	AddSeparatorMessage() tea.Cmd
 	AddOrUpdateToolCall(agentName string, toolCall tools.ToolCall, status types.ToolStatus) tea.Cmd
-	AddToolResult(toolCall tools.ToolCall, result string, status types.ToolStatus) tea.Cmd
-	AppendToLastMessage(agentName string, content string) tea.Cmd
+	AddToolResult(msg *runtime.ToolCallResponseEvent, status types.ToolStatus) tea.Cmd
+	AppendToLastMessage(agentName string, messageType types.MessageType, content string) tea.Cmd
 	ClearMessages()
 	ScrollToBottom() tea.Cmd
 	AddShellOutputMessage(content string) tea.Cmd
@@ -477,7 +478,7 @@ func (m *model) AddShellOutputMessage(content string) tea.Cmd {
 // AddAssistantMessage adds an assistant message to the chat
 func (m *model) AddAssistantMessage() tea.Cmd {
 	return m.addMessage(&types.Message{
-		Type: types.MessageTypeAssistant,
+		Type: types.MessageTypeSpinner,
 	})
 }
 
@@ -505,7 +506,7 @@ func (m *model) addMessage(msg *types.Message) tea.Cmd {
 
 // AddSeparatorMessage adds a separator message to the chat
 func (m *model) AddSeparatorMessage() tea.Cmd {
-	m.removeLastEmptyAssistantMessage()
+	m.removeSpinner()
 	msg := types.Message{
 		Type: types.MessageTypeSeparator,
 	}
@@ -535,7 +536,7 @@ func (m *model) AddOrUpdateToolCall(agentName string, toolCall tools.ToolCall, s
 	}
 
 	// If not found by ID, remove last empty assistant message
-	m.removeLastEmptyAssistantMessage()
+	m.removeSpinner()
 
 	// Create new tool call
 	msg := types.Message{
@@ -553,14 +554,14 @@ func (m *model) AddOrUpdateToolCall(agentName string, toolCall tools.ToolCall, s
 }
 
 // AddToolResult adds tool result to the most recent matching tool call
-func (m *model) AddToolResult(toolCall tools.ToolCall, result string, status types.ToolStatus) tea.Cmd {
+func (m *model) AddToolResult(msg *runtime.ToolCallResponseEvent, status types.ToolStatus) tea.Cmd {
 	for i := len(m.messages) - 1; i >= 0; i-- {
-		msg := &m.messages[i]
-		if msg.ToolCall.ID == toolCall.ID {
-			msg.Content = result
-			msg.ToolStatus = status
+		toolMessage := &m.messages[i]
+		if toolMessage.ToolCall.ID == msg.ToolCall.ID {
+			toolMessage.Content = msg.Response
+			toolMessage.ToolStatus = status
 			// Update the corresponding view
-			view := m.createToolCallView(msg)
+			view := m.createToolCallView(toolMessage)
 			m.views[i] = view
 			return view.Init()
 		}
@@ -569,14 +570,16 @@ func (m *model) AddToolResult(toolCall tools.ToolCall, result string, status typ
 }
 
 // AppendToLastMessage appends content to the last message (for streaming)
-func (m *model) AppendToLastMessage(agentName, content string) tea.Cmd {
+func (m *model) AppendToLastMessage(agentName string, messageType types.MessageType, content string) tea.Cmd {
+	m.removeSpinner()
+
 	if len(m.messages) == 0 {
 		return nil
 	}
 	lastIdx := len(m.messages) - 1
 	lastMsg := &m.messages[lastIdx]
 
-	if lastMsg.Type == types.MessageTypeAssistant {
+	if lastMsg.Type == messageType {
 		wasAtBottom := m.isAtBottom()
 		lastMsg.Content += content
 		lastMsg.Sender = agentName
@@ -599,7 +602,7 @@ func (m *model) AppendToLastMessage(agentName, content string) tea.Cmd {
 	} else {
 		// Create new assistant message
 		msg := types.Message{
-			Type:    types.MessageTypeAssistant,
+			Type:    messageType,
 			Content: content,
 			Sender:  agentName,
 		}
@@ -653,13 +656,13 @@ func (m *model) createMessageView(msg *types.Message) layout.Model {
 	return view
 }
 
-// removeLastEmptyAssistantMessage removes the last message if it's an assistant message with empty content
-func (m *model) removeLastEmptyAssistantMessage() {
+// removeSpinner removes the last message if it's an assistant message with empty content
+func (m *model) removeSpinner() {
 	if len(m.messages) > 0 {
 		lastIdx := len(m.messages) - 1
 		lastMessage := m.messages[lastIdx]
 
-		if lastMessage.Type == types.MessageTypeAssistant && strings.Trim(lastMessage.Content, "\r\n\t ") == "" {
+		if lastMessage.Type == types.MessageTypeSpinner {
 			m.messages = m.messages[:lastIdx]
 			if len(m.views) > lastIdx {
 				m.views = m.views[:lastIdx]
diff --git a/internal/tui/page/chat/chat.go b/internal/tui/page/chat/chat.go
@@ -65,6 +65,7 @@ type KeyMap struct {
 	Tab    key.Binding
 	Quit   key.Binding
 	Cancel key.Binding
+	Copy   key.Binding
 }
 
 // DefaultKeyMap returns the default key bindings
@@ -203,7 +204,10 @@ func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		cmd := p.messages.AddAssistantMessage()
 		return p, tea.Batch(cmd, p.messages.ScrollToBottom(), spinnerCmd)
 	case *runtime.AgentChoiceEvent:
-		cmd := p.messages.AppendToLastMessage(msg.AgentName, msg.Choice.Delta.Content)
+		cmd := p.messages.AppendToLastMessage(msg.AgentName, types.MessageTypeAssistant, msg.Content)
+		return p, tea.Batch(cmd, p.messages.ScrollToBottom())
+	case *runtime.AgentChoiceReasoningEvent:
+		cmd := p.messages.AppendToLastMessage(msg.AgentName, types.MessageTypeAssistantReasoning, msg.Content)
 		return p, tea.Batch(cmd, p.messages.ScrollToBottom())
 	case *runtime.SessionTitleEvent:
 		p.sessionTitle = msg.Title
@@ -249,7 +253,7 @@ func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		return p, tea.Batch(cmd, p.messages.ScrollToBottom(), spinnerCmd)
 	case *runtime.ToolCallResponseEvent:
 		spinnerCmd := p.setWorking(true)
-		cmd := p.messages.AddToolResult(msg.ToolCall, msg.Response, types.ToolStatusCompleted)
+		cmd := p.messages.AddToolResult(msg, types.ToolStatusCompleted)
 		return p, tea.Batch(cmd, p.messages.ScrollToBottom(), spinnerCmd)
 	}
 
diff --git a/internal/tui/types/types.go b/internal/tui/types/types.go
@@ -8,6 +8,8 @@ type MessageType int
 const (
 	MessageTypeUser MessageType = iota
 	MessageTypeAssistant
+	MessageTypeAssistantReasoning
+	MessageTypeSpinner
 	MessageTypeError
 	MessageTypeShellOutput
 	MessageTypeSeparator
diff --git a/pkg/chat/chat.go b/pkg/chat/chat.go
@@ -86,10 +86,11 @@ const (
 
 // MessageDelta represents a delta/chunk in a streaming response
 type MessageDelta struct {
-	Role         string              `json:"role,omitempty"`
-	Content      string              `json:"content,omitempty"`
-	FunctionCall *tools.FunctionCall `json:"function_call,omitempty"`
-	ToolCalls    []tools.ToolCall    `json:"tool_calls,omitempty"`
+	Role             string              `json:"role,omitempty"`
+	Content          string              `json:"content,omitempty"`
+	ReasoningContent string              `json:"reasoning_content,omitempty"`
+	FunctionCall     *tools.FunctionCall `json:"function_call,omitempty"`
+	ToolCalls        []tools.ToolCall    `json:"tool_calls,omitempty"`
 }
 
 // MessageStreamChoice represents a choice in a streaming response
diff --git a/pkg/model/provider/anthropic/adapter.go b/pkg/model/provider/anthropic/adapter.go
@@ -68,7 +68,8 @@ func (a *streamAdapter) Recv() (chat.MessageStreamResponse, error) {
 		switch deltaVariant := eventVariant.Delta.AsAny().(type) {
 		case anthropic.TextDelta:
 			response.Choices[0].Delta.Content = deltaVariant.Text
-
+		case anthropic.ThinkingDelta:
+			response.Choices[0].Delta.ReasoningContent = deltaVariant.Thinking
 		case anthropic.InputJSONDelta:
 			inputBytes := deltaVariant.PartialJSON
 			toolCall := tools.ToolCall{
diff --git a/pkg/model/provider/oaistream/adapter.go b/pkg/model/provider/oaistream/adapter.go
@@ -78,8 +78,9 @@ func (a *StreamAdapter) Recv() (chat.MessageStreamResponse, error) {
 			Index:        choice.Index,
 			FinishReason: finishReason,
 			Delta: chat.MessageDelta{
-				Role:    choice.Delta.Role,
-				Content: choice.Delta.Content,
+				Role:             choice.Delta.Role,
+				Content:          choice.Delta.Content,
+				ReasoningContent: choice.Delta.ReasoningContent,
 			},
 		}
 
diff --git a/pkg/runtime/event.go b/pkg/runtime/event.go
@@ -1,7 +1,6 @@
 package runtime
 
 import (
-	"github.com/docker/cagent/pkg/chat"
 	"github.com/docker/cagent/pkg/tools"
 )
 
@@ -120,20 +119,35 @@ func (e *StreamStartedEvent) GetAgentName() string {
 func (e *StreamStartedEvent) isEvent() {}
 
 type AgentChoiceEvent struct {
-	Type   string                   `json:"type"`
-	Choice chat.MessageStreamChoice `json:"choice"`
+	Type    string `json:"type"`
+	Content string `json:"content"`
 	AgentContext
 }
 
-func AgentChoice(agentName string, choice chat.MessageStreamChoice) Event { //nolint:gocritic
+func AgentChoice(agentName string, content string) Event { //nolint:gocritic
 	return &AgentChoiceEvent{
 		Type:         "agent_choice",
-		Choice:       choice,
+		Content:      content,
 		AgentContext: AgentContext{AgentName: agentName},
 	}
 }
 func (e *AgentChoiceEvent) isEvent() {}
 
+type AgentChoiceReasoningEvent struct {
+	Type    string `json:"type"`
+	Content string `json:"content"`
+	AgentContext
+}
+
+func AgentChoiceReasoning(agentName string, content string) Event { //nolint:gocritic
+	return &AgentChoiceReasoningEvent{
+		Type:         "agent_choice_reasoning",
+		Content:      content,
+		AgentContext: AgentContext{AgentName: agentName},
+	}
+}
+func (e *AgentChoiceReasoningEvent) isEvent() {}
+
 type ErrorEvent struct {
 	Type  string `json:"type"`
 	Error string `json:"error"`
diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
@@ -209,7 +209,7 @@ func (r *Runtime) RunStream(ctx context.Context, sess *session.Session) <-chan E
 			}
 
 			slog.Debug("Processing stream", "agent", a.Name())
-			calls, content, stopped, err := r.handleStream(ctx, stream, a, sess, m, events)
+			calls, content, reasoningContent, stopped, err := r.handleStream(ctx, stream, a, sess, m, events)
 			if err != nil {
 				// Treat context cancellation as a graceful stop
 				if errors.Is(err, context.Canceled) {
@@ -238,10 +238,11 @@ func (r *Runtime) RunStream(ctx context.Context, sess *session.Session) <-chan E
 
 			// Add assistant message to conversation history
 			assistantMessage := chat.Message{
-				Role:      chat.MessageRoleAssistant,
-				Content:   content,
-				ToolCalls: calls,
-				CreatedAt: time.Now().Format(time.RFC3339),
+				Role:             chat.MessageRoleAssistant,
+				Content:          content,
+				ReasoningContent: reasoningContent,
+				ToolCalls:        calls,
+				CreatedAt:        time.Now().Format(time.RFC3339),
 			}
 
 			sess.AddMessage(session.NewAgentMessage(a, &assistantMessage))
@@ -317,10 +318,11 @@ func (r *Runtime) Run(ctx context.Context, sess *session.Session) ([]session.Mes
 }
 
 // handleStream handles the stream processing
-func (r *Runtime) handleStream(ctx context.Context, stream chat.MessageStream, a *agent.Agent, sess *session.Session, m *modelsdev.Model, events chan Event) (calls []tools.ToolCall, content string, stopped bool, err error) {
+func (r *Runtime) handleStream(ctx context.Context, stream chat.MessageStream, a *agent.Agent, sess *session.Session, m *modelsdev.Model, events chan Event) (calls []tools.ToolCall, content, reasoningContent string, stopped bool, err error) {
 	defer stream.Close()
 
 	var fullContent strings.Builder
+	var fullReasoningContent strings.Builder
 	var toolCalls []tools.ToolCall
 	// Track which tool call indices we've already emitted partial events for
 	emittedPartialEvents := make(map[string]bool)
@@ -331,7 +333,7 @@ func (r *Runtime) handleStream(ctx context.Context, stream chat.MessageStream, a
 			break
 		}
 		if err != nil {
-			return nil, "", true, fmt.Errorf("error receiving from stream: %w", err)
+			return nil, "", "", true, fmt.Errorf("error receiving from stream: %w", err)
 		}
 
 		if response.Usage != nil {
@@ -360,7 +362,7 @@ func (r *Runtime) handleStream(ctx context.Context, stream chat.MessageStream, a
 		}
 		choice := response.Choices[0]
 		if choice.FinishReason == chat.FinishReasonStop {
-			return toolCalls, fullContent.String(), true, nil
+			return toolCalls, fullContent.String(), fullReasoningContent.String(), true, nil
 		}
 
 		// Handle tool calls
@@ -415,13 +417,18 @@ func (r *Runtime) handleStream(ctx context.Context, stream chat.MessageStream, a
 			continue
 		}
 
+		if choice.Delta.ReasoningContent != "" {
+			events <- AgentChoiceReasoning(a.Name(), choice.Delta.ReasoningContent)
+			fullReasoningContent.WriteString(choice.Delta.ReasoningContent)
+		}
+
 		if choice.Delta.Content != "" {
-			events <- AgentChoice(a.Name(), choice)
+			events <- AgentChoice(a.Name(), choice.Delta.Content)
 			fullContent.WriteString(choice.Delta.Content)
 		}
 	}
 
-	return toolCalls, fullContent.String(), false, nil
+	return toolCalls, fullContent.String(), fullReasoningContent.String(), false, nil
 }
 
 // processToolCalls handles the execution of tool calls for an agent

Original file line number	Diff line number	Diff line change
`@@ -107,7 +107,7 @@ func NewNewCmd() *cobra.Command {`
`107`	`107`	`fmt.Println()`
`108`	`108`	`llmIsTyping = true`
`109`	`109`	`}`
`110`		`- fmt.Printf("%s", e.Choice.Delta.Content)`
	`110`	`+ fmt.Printf("%s", e.Content)`
`111`	`111`	`case *runtime.ToolCallEvent:`
`112`	`112`	`if llmIsTyping {`
`113`	`113`	`fmt.Println()`
Original file line number	Diff line number	Diff line change
`@@ -321,7 +321,7 @@ func runWithoutTUI(ctx context.Context, agentFilename string, rt *runtime.Runtim`
`321`	`321`	`}`
`322`	`322`	`llmIsTyping = true`
`323`	`323`	`}`
`324`		`- fmt.Printf("%s", e.Choice.Delta.Content)`
	`324`	`+ fmt.Printf("%s", e.Content)`
`325`	`325`	`case *runtime.ToolCallConfirmationEvent:`
`326`	`326`	`if llmIsTyping {`
`327`	`327`	`fmt.Println()`