From f0bdd90aa2459fa59e5f91acf427109e381c38e7 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Tue, 4 Aug 2026 15:43:15 +0800 Subject: [PATCH] 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 | 53 +++- .../agent/component/agent_artifact_test.go | 5 +- internal/agent/component/agent_test.go | 269 ++++++++++++++++++ internal/agent/tool/code_exec.go | 14 +- internal/agent/tool/code_exec_test.go | 62 ++++ 5 files changed, 395 insertions(+), 8 deletions(-) diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index 1ba3fe3e0a..b36d85d23d 100644 --- a/internal/agent/component/agent.go +++ b/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 diff --git a/internal/agent/component/agent_artifact_test.go b/internal/agent/component/agent_artifact_test.go index f84cbf7be0..e6554cb252 100644 --- a/internal/agent/component/agent_artifact_test.go +++ b/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 } diff --git a/internal/agent/component/agent_test.go b/internal/agent/component/agent_test.go index ea89d21659..642ba8c61f 100644 --- a/internal/agent/component/agent_test.go +++ b/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{ diff --git a/internal/agent/tool/code_exec.go b/internal/agent/tool/code_exec.go index 02f2365c3d..ec883be04a 100644 --- a/internal/agent/tool/code_exec.go +++ b/internal/agent/tool/code_exec.go @@ -43,7 +43,9 @@ var ErrCodeExecSandboxMissing = errors.New( const codeExecToolName = "execute_code" const codeExecToolDescription = "This tool has a sandbox that can execute code written in 'Python'/'Javascript'. " + - "It receives a piece of code and returns a JSON string." + "It receives a piece of code and returns a JSON string. " + + "The code must define a main function (Python) or export main (JavaScript); " + + "the return value of main is returned as the tool result." // codeExecArgs is the JSON shape the model sends in. The Python // tool accepts "lang" + "script"; we also accept "code" as a @@ -100,15 +102,15 @@ func (c *CodeExecTool) Info(_ context.Context) (*schema.ToolInfo, error) { Name: codeExecToolName, Desc: codeExecToolDescription, ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ - "language": { + "lang": { Type: schema.String, - Desc: "The programming language of the code. Allowed: 'python' (or 'python3'), 'javascript' (or 'nodejs').", - Enum: []string{"python", "python3", "javascript", "nodejs"}, + Desc: "The programming language of this piece of code.", + Enum: []string{"python", "javascript"}, Required: true, }, - "code": { + "script": { Type: schema.String, - Desc: "The code to execute. Must define a `main` function (Python) or export `main` (JavaScript).", + Desc: "A piece of code in the correct format. It must define main(...).", Required: true, }, }), diff --git a/internal/agent/tool/code_exec_test.go b/internal/agent/tool/code_exec_test.go index 7948944eee..369d34378a 100644 --- a/internal/agent/tool/code_exec_test.go +++ b/internal/agent/tool/code_exec_test.go @@ -91,6 +91,68 @@ func TestCodeExec_Info(t *testing.T) { if !strings.Contains(info.Desc, "Python") { t.Errorf("Desc = %q, want to mention Python", info.Desc) } + + params, err := info.ParamsOneOf.ToJSONSchema() + if err != nil { + t.Fatalf("Info schema: %v", err) + } + encoded, err := json.Marshal(params) + if err != nil { + t.Fatalf("marshal Info schema: %v", err) + } + var schema map[string]any + if err := json.Unmarshal(encoded, &schema); err != nil { + t.Fatalf("decode Info schema: %v", err) + } + properties, ok := schema["properties"].(map[string]any) + if !ok { + t.Fatalf("Info schema properties = %#v, want object", schema["properties"]) + } + for _, name := range []string{"lang", "script"} { + if _, ok := properties[name]; !ok { + t.Errorf("Info schema missing %q", name) + } + } + for _, name := range []string{"language", "code", "arguments", "outputs"} { + if _, ok := properties[name]; ok { + t.Errorf("Info schema unexpectedly exposes node field %q", name) + } + } + required, ok := schema["required"].([]any) + if !ok { + t.Fatalf("Info schema required = %#v, want array", schema["required"]) + } + requiredFields := make(map[string]bool, len(required)) + for _, field := range required { + if name, ok := field.(string); ok { + requiredFields[name] = true + } + } + if !requiredFields["lang"] || !requiredFields["script"] { + t.Errorf("Info schema required = %#v, want lang and script", required) + } + langProp, ok := properties["lang"].(map[string]any) + if !ok { + t.Fatalf("lang property = %#v, want object", properties["lang"]) + } + if typ, _ := langProp["type"].(string); typ != "string" { + t.Errorf("lang.type = %q, want string", typ) + } + enum, ok := langProp["enum"].([]any) + if !ok { + t.Fatalf("lang.enum = %#v, want array", langProp["enum"]) + } + gotEnum := make([]string, len(enum)) + for i, e := range enum { + s, ok := e.(string) + if !ok { + t.Fatalf("lang.enum[%d] = %#v, want string", i, e) + } + gotEnum[i] = s + } + if len(gotEnum) != 2 || gotEnum[0] != "python" || gotEnum[1] != "javascript" { + t.Errorf("lang.enum = %v, want [python javascript]", gotEnum) + } } // TestCodeExec_ResultExtractsArtifacts pins the artifact