From a53e4a4a64c4d914ae4506b20589764c39a49c6f Mon Sep 17 00:00:00 2001 From: tian-IRT <59220904+tian-IRT@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:52:04 +0800 Subject: [PATCH] fix: recover from duplicate streaming tool call indexes (#231) --- internal/multiagent/eino_orchestration.go | 5 +- internal/multiagent/eino_single_runner.go | 3 +- .../multiagent/eino_stream_tool_call_index.go | 117 ++++++++++++++++++ .../eino_stream_tool_call_index_test.go | 117 ++++++++++++++++++ internal/multiagent/runner.go | 12 +- 5 files changed, 246 insertions(+), 8 deletions(-) create mode 100644 internal/multiagent/eino_stream_tool_call_index.go create mode 100644 internal/multiagent/eino_stream_tool_call_index_test.go diff --git a/internal/multiagent/eino_orchestration.go b/internal/multiagent/eino_orchestration.go index 66626b5a..a0ad6829 100644 --- a/internal/multiagent/eino_orchestration.go +++ b/internal/multiagent/eino_orchestration.go @@ -9,7 +9,6 @@ import ( "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" - "github.com/cloudwego/eino-ext/components/model/openai" "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk/prebuilt/planexecute" "github.com/cloudwego/eino/components/model" @@ -19,8 +18,8 @@ import ( // PlanExecuteRootArgs 构建 Eino adk/prebuilt/planexecute 根 Agent 所需参数。 type PlanExecuteRootArgs struct { - MainToolCallingModel *openai.ChatModel - ExecModel *openai.ChatModel + MainToolCallingModel model.ToolCallingChatModel + ExecModel model.ToolCallingChatModel OrchInstruction string ToolsCfg adk.ToolsConfig ExecMaxIter int diff --git a/internal/multiagent/eino_single_runner.go b/internal/multiagent/eino_single_runner.go index ea7d255b..c18b72ed 100644 --- a/internal/multiagent/eino_single_runner.go +++ b/internal/multiagent/eino_single_runner.go @@ -121,10 +121,11 @@ func RunEinoSingleChatModelAgent( } reasoning.ApplyToEinoChatModelConfig(baseModelCfg, &appCfg.OpenAI, reasoningClient) - mainModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) + baseMainModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) if err != nil { return nil, fmt.Errorf("eino single 模型: %w", err) } + mainModel := newStreamToolCallIndexRepairModel(baseMainModel) mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) if err != nil { diff --git a/internal/multiagent/eino_stream_tool_call_index.go b/internal/multiagent/eino_stream_tool_call_index.go new file mode 100644 index 00000000..2d769f7a --- /dev/null +++ b/internal/multiagent/eino_stream_tool_call_index.go @@ -0,0 +1,117 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// streamToolCallIndexRepairModel isolates an OpenAI-compatible streaming +// protocol defect before Eino concatenates response chunks. Some providers +// reuse a tool-call index for different non-empty tool-call IDs in one stream. +// Eino correctly rejects that shape because one index represents one call. +// +// The wrapper keeps valid streams untouched. When it sees the conflicting +// shape, it assigns each distinct ID a stable, stream-local index so Eino can +// retain all calls instead of aborting the agent run. +type streamToolCallIndexRepairModel struct { + base model.ToolCallingChatModel +} + +func newStreamToolCallIndexRepairModel(base model.ToolCallingChatModel) model.ToolCallingChatModel { + if base == nil { + return nil + } + return &streamToolCallIndexRepairModel{base: base} +} + +func (m *streamToolCallIndexRepairModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return m.base.Generate(ctx, input, opts...) +} + +func (m *streamToolCallIndexRepairModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + stream, err := m.base.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + state := newStreamToolCallIndexRepairState() + return schema.StreamReaderWithConvert(stream, state.repairMessage), nil +} + +func (m *streamToolCallIndexRepairModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + withTools, err := m.base.WithTools(tools) + if err != nil { + return nil, err + } + return newStreamToolCallIndexRepairModel(withTools), nil +} + +type streamToolCallIndexRepairState struct { + indexByID map[string]int + idByIndex map[int]string + nextFreeIndex int +} + +func newStreamToolCallIndexRepairState() *streamToolCallIndexRepairState { + return &streamToolCallIndexRepairState{ + indexByID: make(map[string]int), + idByIndex: make(map[int]string), + } +} + +func (s *streamToolCallIndexRepairState) repairMessage(msg *schema.Message) (*schema.Message, error) { + if msg == nil || len(msg.ToolCalls) == 0 { + return msg, nil + } + + var calls []schema.ToolCall + changed := false + for i := range msg.ToolCalls { + call := msg.ToolCalls[i] + if call.Index == nil || call.ID == "" { + continue + } + + sourceIndex := *call.Index + if sourceIndex >= s.nextFreeIndex { + s.nextFreeIndex = sourceIndex + 1 + } + + assigned, known := s.indexByID[call.ID] + if !known { + assigned = sourceIndex + if owner, occupied := s.idByIndex[assigned]; occupied && owner != call.ID { + assigned = s.takeFreeIndex() + } + s.indexByID[call.ID] = assigned + s.idByIndex[assigned] = call.ID + } + if assigned == sourceIndex { + continue + } + if calls == nil { + calls = append([]schema.ToolCall(nil), msg.ToolCalls...) + } + index := assigned + calls[i].Index = &index + changed = true + } + + if !changed { + return msg, nil + } + out := *msg + out.ToolCalls = calls + return &out, nil +} + +func (s *streamToolCallIndexRepairState) takeFreeIndex() int { + for { + candidate := s.nextFreeIndex + s.nextFreeIndex++ + if _, occupied := s.idByIndex[candidate]; !occupied { + return candidate + } + } +} diff --git a/internal/multiagent/eino_stream_tool_call_index_test.go b/internal/multiagent/eino_stream_tool_call_index_test.go new file mode 100644 index 00000000..abcdfb90 --- /dev/null +++ b/internal/multiagent/eino_stream_tool_call_index_test.go @@ -0,0 +1,117 @@ +package multiagent + +import ( + "context" + "io" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type streamToolCallIndexFakeModel struct { + chunks []*schema.Message +} + +func (m *streamToolCallIndexFakeModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, nil +} + +func (m *streamToolCallIndexFakeModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray(m.chunks), nil +} + +func (m *streamToolCallIndexFakeModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func TestStreamToolCallIndexRepairSeparatesConflictingIDs(t *testing.T) { + index := 0 + wrapped := newStreamToolCallIndexRepairModel(&streamToolCallIndexFakeModel{chunks: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{{ + Index: &index, ID: "fc_call_0", Type: "function", + Function: schema.FunctionCall{Name: "search", Arguments: `{"query":"one"}`}, + }}), + schema.AssistantMessage("", []schema.ToolCall{{ + Index: &index, ID: "fc_call_1", Type: "function", + Function: schema.FunctionCall{Name: "task", Arguments: `{"query":"two"}`}, + }}), + }}) + + got := readStreamToolCallChunks(t, wrapped) + merged, err := schema.ConcatMessages(got) + if err != nil { + t.Fatalf("ConcatMessages() error = %v", err) + } + if len(merged.ToolCalls) != 2 { + t.Fatalf("tool call count = %d, want 2", len(merged.ToolCalls)) + } + if merged.ToolCalls[0].ID != "fc_call_0" || merged.ToolCalls[1].ID != "fc_call_1" { + t.Fatalf("tool call IDs = %#v", merged.ToolCalls) + } + if merged.ToolCalls[0].Index == nil || *merged.ToolCalls[0].Index != 0 || merged.ToolCalls[1].Index == nil || *merged.ToolCalls[1].Index != 1 { + t.Fatalf("tool call indexes = %#v", merged.ToolCalls) + } +} + +func TestStreamToolCallIndexRepairPreservesFragmentsForOneID(t *testing.T) { + index := 0 + wrapped := newStreamToolCallIndexRepairModel(&streamToolCallIndexFakeModel{chunks: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{{ + Index: &index, ID: "call_0", Type: "function", + Function: schema.FunctionCall{Name: "search", Arguments: `{"query":"`}, + }}), + schema.AssistantMessage("", []schema.ToolCall{{ + Index: &index, ID: "call_0", Type: "function", + Function: schema.FunctionCall{Arguments: `one"}`}, + }}), + }}) + + got := readStreamToolCallChunks(t, wrapped) + merged, err := schema.ConcatMessages(got) + if err != nil { + t.Fatalf("ConcatMessages() error = %v", err) + } + if len(merged.ToolCalls) != 1 || merged.ToolCalls[0].Function.Arguments != `{"query":"one"}` { + t.Fatalf("tool calls = %#v", merged.ToolCalls) + } +} + +func TestStreamToolCallIndexRepairLeavesValidParallelIndexesUntouched(t *testing.T) { + first, second := 0, 1 + wrapped := newStreamToolCallIndexRepairModel(&streamToolCallIndexFakeModel{chunks: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + {Index: &first, ID: "call_0", Type: "function", Function: schema.FunctionCall{Name: "search", Arguments: `{}`}}, + {Index: &second, ID: "call_1", Type: "function", Function: schema.FunctionCall{Name: "task", Arguments: `{}`}}, + }), + }}) + + got := readStreamToolCallChunks(t, wrapped) + if len(got) != 1 || len(got[0].ToolCalls) != 2 { + t.Fatalf("chunks = %#v", got) + } + if *got[0].ToolCalls[0].Index != 0 || *got[0].ToolCalls[1].Index != 1 { + t.Fatalf("tool call indexes changed: %#v", got[0].ToolCalls) + } +} + +func readStreamToolCallChunks(t *testing.T, chatModel model.ToolCallingChatModel) []*schema.Message { + t.Helper() + stream, err := chatModel.Stream(context.Background(), nil) + if err != nil { + t.Fatalf("Stream() error = %v", err) + } + defer stream.Close() + + var chunks []*schema.Message + for { + chunk, recvErr := stream.Recv() + if recvErr == io.EOF { + return chunks + } + if recvErr != nil { + t.Fatalf("Recv() error = %v", recvErr) + } + chunks = append(chunks, chunk) + } +} diff --git a/internal/multiagent/runner.go b/internal/multiagent/runner.go index 36a29264..fb919f82 100644 --- a/internal/multiagent/runner.go +++ b/internal/multiagent/runner.go @@ -218,10 +218,11 @@ func RunDeepAgent( } } - subModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) + baseSubModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) if err != nil { return nil, fmt.Errorf("子代理 %q ChatModel: %w", id, err) } + subModel := newStreamToolCallIndexRepairModel(baseSubModel) subDefs := ag.ToolsForRole(roleTools) subTools, err := einomcp.ToolsFromDefinitions(ag, holder, subDefs, recorder, nil, toolInvokeNotify, id) @@ -308,10 +309,11 @@ func RunDeepAgent( } } - mainModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) + baseMainModel, err := einoopenai.NewChatModel(ctx, baseModelCfg) if err != nil { return nil, fmt.Errorf("多代理主模型: %w", err) } + mainModel := newStreamToolCallIndexRepairModel(baseMainModel) mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) if err != nil { @@ -481,19 +483,21 @@ func RunDeepAgent( MaxCompletionTokens: &maxCompletionTokens, } reasoning.ApplyPlanExecutePlannerModelConfig(plannerModelCfg, &appCfg.OpenAI) - peMainModel, perr := einoopenai.NewChatModel(ctx, plannerModelCfg) + basePEMainModel, perr := einoopenai.NewChatModel(ctx, plannerModelCfg) if perr != nil { return nil, fmt.Errorf("plan_execute 规划模型: %w", perr) } + peMainModel := newStreamToolCallIndexRepairModel(basePEMainModel) if logger != nil { logger.Info("plan_execute: planner/replanner 使用无 reasoning 的独立 ChatModel（ToolChoiceForced 兼容）", zap.String("model", appCfg.OpenAI.Model), ) } - execModel, perr := einoopenai.NewChatModel(ctx, baseModelCfg) + baseExecModel, perr := einoopenai.NewChatModel(ctx, baseModelCfg) if perr != nil { return nil, fmt.Errorf("plan_execute 执行器模型: %w", perr) } + execModel := newStreamToolCallIndexRepairModel(baseExecModel) // 构建 filesystem 中间件（与 Deep sub-agent 一致） var peFsMw adk.ChatModelAgentMiddleware if einoSkillMW != nil && einoFSTools && einoLoc != nil {