Merge pull request #1322 from trheyi/main

Enhance streaming message handling in DefaultStreamHandler

agent/llm/handlers/stream.go

@@ -62,10 +62,11 @@ func DefaultStreamHandler(ctx *context.Context) context.StreamFunc {
 
 // streamState manages the state of the streaming process
 type streamState struct {
-	ctx       *context.Context
+	ctx         *context.Context
-	inGroup   bool
+	inGroup     bool
-	currentID string
+	currentID   string
-	buffer    []byte
+	currentType string // Track the current message type (text, thinking, tool_call)
+	buffer      []byte
 }
 
 // handleStreamStart handles stream start event
@@ -96,6 +97,9 @@ func (s *streamState) handleText(data []byte) int {
 		s.currentID = generateMessageID()
 	}
 
+	// Track current message type
+	s.currentType = message.TypeText
+
 	// Append to buffer
 	s.buffer = append(s.buffer, data...)
 
@@ -128,6 +132,9 @@ func (s *streamState) handleThinking(data []byte) int {
 		s.currentID = generateMessageID()
 	}
 
+	// Track current message type
+	s.currentType = message.TypeThinking
+
 	// Append to buffer
 	s.buffer = append(s.buffer, data...)
 
@@ -190,9 +197,15 @@ func (s *streamState) handleGroupEnd(data []byte) int {
 
 	// Send done message with complete content
 	if s.currentID != "" && len(s.buffer) > 0 {
+		// Use the tracked message type (thinking, text, tool_call, etc.)
+		msgType := s.currentType
+		if msgType == "" {
+			msgType = message.TypeText // Fallback to text if type not set
+		}
+
 		msg := &message.Message{
 			ID:   s.currentID,
-			Type: message.TypeText, // Default to text
+			Type: msgType, // Use the actual message type from the group
 			Done: true,
 			Props: map[string]interface{}{
 				"content": string(s.buffer),
@@ -204,6 +217,7 @@ func (s *streamState) handleGroupEnd(data []byte) int {
 	// Reset state
 	s.inGroup = false
 	s.currentID = ""
+	s.currentType = ""
 	s.buffer = []byte{}
 
 	return 0 // Continue

agent/llm/providers/openai/deepseek_r1_test.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 	"testing"
 
+	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
@@ -59,20 +60,58 @@ func TestDeepSeekR1StreamBasic(t *testing.T) {
 	// Create context
 	ctx := newDeepSeekTestContext("test-deepseek-r1-basic", "deepseek.r1")
 
-	// Track streaming chunks
+	// Track streaming chunks and group events
 	var reasoningChunks []string
 	var contentChunks []string
+	var thinkingGroupEnded bool
+	var textGroupEnded bool
+
 	handler := func(chunkType context.StreamChunkType, data []byte) int {
 		dataStr := string(data)
 		t.Logf("Stream chunk [%s]: %s", chunkType, dataStr)
 
 		// Track different chunk types
-		if chunkType == context.ChunkThinking {
+		switch chunkType {
+		case context.ChunkThinking:
 			reasoningChunks = append(reasoningChunks, dataStr)
-		} else if chunkType == context.ChunkText {
+		case context.ChunkText:
 			contentChunks = append(contentChunks, dataStr)
 		}
 
+		// Track group_end events to verify type field
+		if chunkType == context.ChunkGroupEnd {
+			// Parse the group_end data to check the type field
+			var groupEndData struct {
+				GroupID    string `json:"group_id"`
+				Type       string `json:"type"`
+				Timestamp  int64  `json:"timestamp"`
+				DurationMs int64  `json:"duration_ms"`
+				ChunkCount int    `json:"chunk_count"`
+				Status     string `json:"status"`
+			}
+
+			if err := jsoniter.Unmarshal(data, &groupEndData); err == nil {
+				t.Logf("✓ group_end received: type=%s, chunks=%d, duration=%dms",
+					groupEndData.Type, groupEndData.ChunkCount, groupEndData.DurationMs)
+
+				// Verify the type field matches expected group types
+				switch groupEndData.Type {
+				case "thinking":
+					thinkingGroupEnded = true
+					if groupEndData.ChunkCount == 0 {
+						t.Error("thinking group_end should have chunk_count > 0")
+					}
+				case "text":
+					textGroupEnded = true
+					if groupEndData.ChunkCount == 0 {
+						t.Error("text group_end should have chunk_count > 0")
+					}
+				}
+			} else {
+				t.Errorf("Failed to parse group_end data: %v", err)
+			}
+		}
+
 		return 0 // Continue
 	}
 
@@ -139,6 +178,19 @@ func TestDeepSeekR1StreamBasic(t *testing.T) {
 		t.Logf("Received %d content chunks", len(contentChunks))
 	}
 
+	// Verify group_end events were received with correct types
+	if !thinkingGroupEnded {
+		t.Error("❌ Expected thinking group_end event but didn't receive it")
+	} else {
+		t.Log("✅ Thinking group_end event received with type='thinking'")
+	}
+
+	if !textGroupEnded {
+		t.Error("❌ Expected text group_end event but didn't receive it")
+	} else {
+		t.Log("✅ Text group_end event received with type='text'")
+	}
+
 	t.Logf("Final response: %+v", response)
 }