fix(go-agent): preserve realtime stream deltas (#17791)

- Start the Agent model-stream collector before ReAct execution.
- Preserve all streamed reasoning/content deltas and drain the collector
on errors.
- Add coverage for delayed thinking streams and tool-call execution.

internal/agent/component/agent.go

@@ -172,6 +172,11 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro
 		ToolsConfig: compose.ToolsNodeConfig{
 			Tools: tools,
 		},
+		// Python's streaming tool loop consumes the complete provider
+		// response before deciding whether the round contains a tool call.
+		// Eino's default checker only inspects the first non-empty chunk,
+		// which can miss a ToolCall emitted after explanatory text.
+		StreamToolCallChecker: scanAllStreamForToolCall,
 		MessageModifier: func(_ context.Context, msgs []*schema.Message) []*schema.Message {
 			if p.SystemPrompt != "" {
 				return append([]*schema.Message{schema.SystemMessage(p.SystemPrompt)}, msgs...)
@@ -186,12 +191,24 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro
 
 	opt, future := react.WithMessageFuture()
 	ctx = setArtifactCollector(ctx, future)
+	// Start the model-stream collector BEFORE agent.Stream. The checker
+	// (scanAllStreamForToolCall) must consume the whole round before the
+	// graph releases its output stream, so agent.Stream does not return
+	// until the model finishes. GetMessageStreams blocks on the future's
+	// started signal (closed by the graph onStart callback), so starting
+	// the collector first lets thinking deltas stream out in real time
+	// while the checker runs, instead of buffering the entire round.
+	emitDone := emitAgentModelStreams(ctx, future)
 	stream, err := agent.Stream(ctx, input, opt)
 	if err != nil {
+		// Drain the collector so its goroutine exits before we return.
+		select {
+		case <-emitDone:
+		case <-ctx.Done():
+		}
 		return nil, err
 	}
 	defer stream.Close()
-	emitDone := emitAgentModelStreams(ctx, future)
 
 	chunks := make([]*schema.Message, 0)
 	for {
@@ -200,6 +217,10 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro
 			break
 		}
 		if err != nil {
+			select {
+			case <-emitDone:
+			case <-ctx.Done():
+			}
 			return nil, err
 		}
 		if chunk == nil {
@@ -220,6 +241,36 @@ func runEinoReActAgent(ctx context.Context, p AgentParam) (*schema.Message, erro
 	return msg, nil
 }
 
+// scanAllStreamForToolCall consumes the whole model response, then branches
+// to the Tools node only when any streamed message contains a ToolCall. It
+// must read to EOF because providers append the tool-call message at the end
+// of the stream (see EinoChatModel.Stream), mirroring Python's
+// async_chat_streamly_with_tools, which likewise consumes an entire SSE round
+// before deciding. This keeps tool_choice=auto — a model may still answer
+// directly when it decides no tool is needed.
+//
+// The checker runs synchronously inside the graph's main loop, so
+// runEinoReActAgent starts emitAgentModelStreams before agent.Stream to keep
+// thinking deltas streaming in real time while this function drains the
+// round.
+func scanAllStreamForToolCall(_ context.Context, stream *schema.StreamReader[*schema.Message]) (bool, error) {
+	defer stream.Close()
+
+	hasToolCall := false
+	for {
+		msg, err := stream.Recv()
+		if err == io.EOF {
+			return hasToolCall, nil
+		}
+		if err != nil {
+			return false, err
+		}
+		if msg != nil && len(msg.ToolCalls) > 0 {
+			hasToolCall = true
+		}
+	}
+}
+
 // buildAgentInputMessages assembles the Python-compatible Agent prompt: the
 // configured history window followed by the current user prompt. The current
 // in-flight user entry is excluded through SnapshotPriorHistory, because the

internal/agent/component/agent_artifact_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/cloudwego/eino/components/model"
@@ -231,7 +232,8 @@ func (m *replayModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChat
 
 // passthroughTool echoes its input as a tool result.
 type passthroughTool struct {
-	name string
+	name  string
+	calls atomic.Int32
 }
 
 func (t *passthroughTool) Info(_ context.Context) (*schema.ToolInfo, error) {
@@ -245,5 +247,6 @@ func (t *passthroughTool) Info(_ context.Context) (*schema.ToolInfo, error) {
 }
 
 func (t *passthroughTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) {
+	t.calls.Add(1)
 	return argumentsInJSON, nil
 }

internal/agent/component/agent_test.go

@@ -20,6 +20,7 @@ import (
 	"io"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/cloudwego/eino/components/model"
@@ -70,6 +71,274 @@ func TestAgent_NoToolsReAct(t *testing.T) {
 	}
 }
 
+func TestScanAllStreamForToolCallWaitsPastTextChunks(t *testing.T) {
+	stream := schema.StreamReaderFromArray([]*schema.Message{
+		{Role: schema.Assistant, Content: "I will calculate this."},
+		{
+			Role: schema.Assistant,
+			ToolCalls: []schema.ToolCall{{
+				ID:   "call-1",
+				Type: "function",
+				Function: schema.FunctionCall{
+					Name:      "execute_code",
+					Arguments: `{"lang":"python","script":"def main(): return 2"}`,
+				},
+			}},
+		},
+	})
+
+	hasToolCall, err := scanAllStreamForToolCall(t.Context(), stream)
+	if err != nil {
+		t.Fatalf("scanAllStreamForToolCall: %v", err)
+	}
+	if !hasToolCall {
+		t.Fatal("scanAllStreamForToolCall = false, want true")
+	}
+}
+
+func TestScanAllStreamForToolCallAllowsDirectAnswer(t *testing.T) {
+	stream := schema.StreamReaderFromArray([]*schema.Message{
+		{Role: schema.Assistant, Content: "No tool is needed."},
+	})
+
+	hasToolCall, err := scanAllStreamForToolCall(t.Context(), stream)
+	if err != nil {
+		t.Fatalf("scanAllStreamForToolCall: %v", err)
+	}
+	if hasToolCall {
+		t.Fatal("scanAllStreamForToolCall = true, want false")
+	}
+}
+
+type textThenToolCallModel struct {
+	turn       int
+	boundTools []*schema.ToolInfo
+}
+
+func (m *textThenToolCallModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
+	m.boundTools = append([]*schema.ToolInfo(nil), tools...)
+	return m, nil
+}
+
+func (m *textThenToolCallModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
+	m.turn++
+	if m.turn == 1 {
+		return &schema.Message{
+			Role: schema.Assistant,
+			ToolCalls: []schema.ToolCall{{
+				ID:   "call-1",
+				Type: "function",
+				Function: schema.FunctionCall{
+					Name:      "execute_code",
+					Arguments: `{"lang":"python","script":"def main(): return 2"}`,
+				},
+			}},
+		}, nil
+	}
+	return &schema.Message{Role: schema.Assistant, Content: "2"}, nil
+}
+
+func (m *textThenToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+	m.turn++
+	if m.turn == 1 {
+		return schema.StreamReaderFromArray([]*schema.Message{
+			{Role: schema.Assistant, Content: "I will calculate this."},
+			{
+				Role: schema.Assistant,
+				ToolCalls: []schema.ToolCall{{
+					ID:   "call-1",
+					Type: "function",
+					Function: schema.FunctionCall{
+						Name:      "execute_code",
+						Arguments: `{"lang":"python","script":"def main(): return 2"}`,
+					},
+				}},
+			},
+		}), nil
+	}
+	return schema.StreamReaderFromArray([]*schema.Message{
+		{Role: schema.Assistant, Content: "2"},
+	}), nil
+}
+
+func TestReactStreamCheckerExecutesToolCallAfterText(t *testing.T) {
+	mdl := &textThenToolCallModel{}
+	tool := &passthroughTool{name: "execute_code"}
+	agent, err := react.NewAgent(t.Context(), &react.AgentConfig{
+		ToolCallingModel: mdl,
+		ToolsConfig: compose.ToolsNodeConfig{
+			Tools: []einotool.BaseTool{tool},
+		},
+		StreamToolCallChecker: scanAllStreamForToolCall,
+		MaxStep:               4,
+	})
+	if err != nil {
+		t.Fatalf("react.NewAgent: %v", err)
+	}
+
+	stream, err := agent.Stream(t.Context(), []*schema.Message{schema.UserMessage("calculate")})
+	if err != nil {
+		t.Fatalf("agent.Stream: %v", err)
+	}
+	defer stream.Close()
+
+	var chunks []*schema.Message
+	for {
+		chunk, recvErr := stream.Recv()
+		if recvErr == io.EOF {
+			break
+		}
+		if recvErr != nil {
+			t.Fatalf("stream.Recv: %v", recvErr)
+		}
+		if chunk != nil {
+			chunks = append(chunks, chunk)
+		}
+	}
+	if mdl.turn < 2 {
+		t.Fatalf("model turns = %d, want a second turn after tool execution", mdl.turn)
+	}
+	if len(mdl.boundTools) != 1 || mdl.boundTools[0].Name != "execute_code" {
+		t.Fatalf("bound tools = %#v, want execute_code", mdl.boundTools)
+	}
+	if got := tool.calls.Load(); got != 1 {
+		t.Fatalf("tool invocations = %d, want exactly 1", got)
+	}
+	final, err := schema.ConcatMessages(chunks)
+	if err != nil {
+		t.Fatalf("schema.ConcatMessages: %v", err)
+	}
+	if final.Content != "2" {
+		t.Fatalf("final content = %q, want 2", final.Content)
+	}
+}
+
+// slowThinkingModel streams a few reasoning chunks (each delayed) before a
+// tool call, then answers directly on the second turn. It is used to prove
+// that thinking deltas reach the agent message emitter while agent.Stream is
+// still blocked inside the stream-tool-call checker.
+type slowThinkingModel struct {
+	turn int
+}
+
+func (m *slowThinkingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
+	return m, nil
+}
+
+func (m *slowThinkingModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
+	m.turn++
+	return &schema.Message{Role: schema.Assistant, Content: "2"}, nil
+}
+
+func (m *slowThinkingModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
+	m.turn++
+	if m.turn == 1 {
+		sr, sw := schema.Pipe[*schema.Message](1)
+		go func() {
+			defer sw.Close()
+			for _, chunk := range []*schema.Message{
+				{Role: schema.Assistant, ReasoningContent: "thinking one"},
+				{Role: schema.Assistant, ReasoningContent: "thinking two"},
+				{Role: schema.Assistant, ReasoningContent: "thinking three"},
+				{
+					Role: schema.Assistant,
+					ToolCalls: []schema.ToolCall{{
+						ID:   "call-1",
+						Type: "function",
+						Function: schema.FunctionCall{
+							Name:      "execute_code",
+							Arguments: `{"lang":"python","script":"def main(): return 2"}`,
+						},
+					}},
+				},
+			} {
+				time.Sleep(30 * time.Millisecond)
+				if sw.Send(chunk, nil) {
+					return
+				}
+			}
+		}()
+		return sr, nil
+	}
+	return schema.StreamReaderFromArray([]*schema.Message{{Role: schema.Assistant, Content: "2"}}), nil
+}
+
+// TestReactCheckerStreamsThinkingBeforeAgentReturns pins the fix for the
+// stream realtime regression: the stream-tool-call checker drains the whole
+// round synchronously inside the graph main loop, so agent.Stream cannot
+// return until the model finishes. The collector (emitAgentModelStreams)
+// must therefore be started before agent.Stream and deliver the first
+// thinking delta while the model is still streaming.
+func TestReactCheckerStreamsThinkingBeforeAgentReturns(t *testing.T) {
+	mdl := &slowThinkingModel{}
+	tool := &passthroughTool{name: "execute_code"}
+	agent, err := react.NewAgent(t.Context(), &react.AgentConfig{
+		ToolCallingModel: mdl,
+		ToolsConfig: compose.ToolsNodeConfig{
+			Tools: []einotool.BaseTool{tool},
+		},
+		StreamToolCallChecker: scanAllStreamForToolCall,
+		MaxStep:               4,
+	})
+	if err != nil {
+		t.Fatalf("react.NewAgent: %v", err)
+	}
+
+	thinkingSeen := make(chan struct{}, 1)
+	ctx := runtime.WithAgentMessageEmitter(t.Context(), func(content, thinking string) {
+		if thinking != "" {
+			select {
+			case thinkingSeen <- struct{}{}:
+			default:
+			}
+		}
+	})
+
+	opt, future := react.WithMessageFuture()
+	emitDone := emitAgentModelStreams(ctx, future)
+
+	streamCh := make(chan *schema.StreamReader[*schema.Message], 1)
+	errCh := make(chan error, 1)
+	go func() {
+		s, streamErr := agent.Stream(ctx, []*schema.Message{schema.UserMessage("calculate")}, opt)
+		if streamErr != nil {
+			errCh <- streamErr
+			return
+		}
+		streamCh <- s
+	}()
+
+	select {
+	case <-thinkingSeen:
+		// First thinking delta arrived while the model is still streaming.
+	case <-streamCh:
+		t.Fatal("agent.Stream returned before any thinking delta was emitted")
+	case <-errCh:
+		t.Fatal("agent.Stream errored before any thinking delta was emitted")
+	case <-time.After(10 * time.Second):
+		t.Fatal("timed out waiting for a thinking delta")
+	}
+
+	var stream *schema.StreamReader[*schema.Message]
+	select {
+	case stream = <-streamCh:
+	case err := <-errCh:
+		t.Fatalf("agent.Stream: %v", err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("agent.Stream did not return")
+	}
+	defer stream.Close()
+	for {
+		if _, recvErr := stream.Recv(); recvErr == io.EOF {
+			break
+		}
+	}
+	if got := tool.calls.Load(); got != 1 {
+		t.Fatalf("tool invocations = %d, want exactly 1", got)
+	}
+	<-emitDone
+}
+
 func TestAgent_EmitsThinking(t *testing.T) {
 	withAgentRunner(t, func(_ context.Context, _ AgentParam) (*schema.Message, error) {
 		return &schema.Message{