diff --git a/internal/multiagent/eino_agentic_stream_block_index_repair.go b/internal/multiagent/eino_agentic_stream_block_index_repair.go new file mode 100644 index 000000000..51da4eb03 --- /dev/null +++ b/internal/multiagent/eino_agentic_stream_block_index_repair.go @@ -0,0 +1,127 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// agenticStreamBlockIndexRepairModel repairs StreamingMeta index collisions on +// the agentic streaming path. Some providers (observed: Cloudflare Workers AI +// gateways) interleave whitespace text chunks between parallel tool-call +// chunks in one SSE stream. The chunk converter in eino-ext acl/openai then +// assigns the text chunk the same StreamingMeta.Index as the preceding +// tool-call block, because it only tracks content-type transitions and does +// not account for tool calls having started. schema.ConcatAgenticMessages +// groups blocks by that index and fails with "content block type mismatch"; +// the ADK retry checker swallows that error and sees a nil message, which +// surfaces as a misleading empty-model-output rejection and an endless +// deterministic retry loop. +// +// The wrapper keeps valid streams untouched. When a block's index was already +// used by a block of a different kind (tool call vs. text/reasoning), the +// block is moved to a fresh index so every concat group stays single-typed. +type agenticStreamBlockIndexRepairModel struct { + base model.AgenticModel +} + +func newAgenticStreamBlockIndexRepairModel(base model.AgenticModel) model.AgenticModel { + if base == nil { + return nil + } + return &agenticStreamBlockIndexRepairModel{base: base} +} + +func (m *agenticStreamBlockIndexRepairModel) Generate( + ctx context.Context, + input []*schema.AgenticMessage, + opts ...model.Option, +) (*schema.AgenticMessage, error) { + return m.base.Generate(ctx, input, opts...) +} + +func (m *agenticStreamBlockIndexRepairModel) Stream( + ctx context.Context, + input []*schema.AgenticMessage, + opts ...model.Option, +) (*schema.StreamReader[*schema.AgenticMessage], error) { + stream, err := m.base.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + state := newAgenticBlockIndexRepairState() + return schema.StreamReaderWithConvert(stream, state.repairMessage), nil +} + +type agenticBlockIndexRepairState struct { + kindByIndex map[int]string // first-seen block kind per original converter index + reassigned map[string]int // "idx|kind" of colliding blocks -> reassigned index + nextFree int // reassigned indices count down from -1 so they can + // never collide with converter-assigned indices (>= 0) that appear later in + // the stream. All chunks of the same colliding (idx, kind) pair share one + // reassigned index, so they still merge into a single block. Reassigned + // blocks sort before index 0 in the final message; block content is fully + // preserved and only the position of stray interleaved text moves. +} + +func newAgenticBlockIndexRepairState() *agenticBlockIndexRepairState { + return &agenticBlockIndexRepairState{ + kindByIndex: make(map[int]string), + reassigned: make(map[string]int), + nextFree: -1, + } +} + +func agenticBlockKind(block *schema.ContentBlock) string { + if block != nil && block.FunctionToolCall != nil { + return "tool" + } + return "content" +} + +func (s *agenticBlockIndexRepairState) repairMessage(msg *schema.AgenticMessage) (*schema.AgenticMessage, error) { + if msg == nil || len(msg.ContentBlocks) == 0 { + return msg, nil + } + var out *schema.AgenticMessage + for i, block := range msg.ContentBlocks { + if block == nil || block.StreamingMeta == nil { + continue + } + idx := block.StreamingMeta.Index + kind := agenticBlockKind(block) + prev, seen := s.kindByIndex[idx] + if !seen { + s.kindByIndex[idx] = kind + continue + } + if prev == kind { + continue + } + // Index collision across block kinds: reassign a fresh (negative) index + // so ConcatAgenticMessages never groups mismatched block types. + key := fmt.Sprintf("%d|%s", idx, kind) + newIdx, ok := s.reassigned[key] + if !ok { + newIdx = s.nextFree + s.nextFree-- + s.reassigned[key] = newIdx + } + if out == nil { + cp := *msg + cp.ContentBlocks = append([]*schema.ContentBlock(nil), msg.ContentBlocks...) + out = &cp + } + newBlock := *block + newMeta := *block.StreamingMeta + newMeta.Index = newIdx + newBlock.StreamingMeta = &newMeta + out.ContentBlocks[i] = &newBlock + } + if out != nil { + return out, nil + } + return msg, nil +} diff --git a/internal/multiagent/eino_agentic_stream_block_index_repair_test.go b/internal/multiagent/eino_agentic_stream_block_index_repair_test.go new file mode 100644 index 000000000..67bd471cc --- /dev/null +++ b/internal/multiagent/eino_agentic_stream_block_index_repair_test.go @@ -0,0 +1,133 @@ +package multiagent + +import ( + "context" + "io" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// agenticBlockIndexFakeModel replays a fixed frame sequence, mimicking a +// provider stream that interleaves a whitespace text chunk between two +// parallel tool-call blocks (observed on Cloudflare Workers AI gateways). +type agenticBlockIndexFakeModel struct { + frames []*schema.AgenticMessage +} + +func (m *agenticBlockIndexFakeModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return schema.ConcatAgenticMessages(m.frames) +} + +func (m *agenticBlockIndexFakeModel) Stream(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + sr, sw := schema.Pipe[*schema.AgenticMessage](len(m.frames) + 1) + go func() { + for _, f := range m.frames { + sw.Send(f, nil) + } + sw.Close() + }() + return sr, nil +} + +func streamBlock(idx int, block *schema.ContentBlock) *schema.AgenticMessage { + block.StreamingMeta = &schema.StreamingMeta{Index: idx} + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{block}, + } +} + +func TestAgenticStreamBlockIndexRepair_InterleavedTextBetweenToolCalls(t *testing.T) { + // Converter output for: text -> tool call 1 -> interleaved "\n" text -> tool call 2. + // The interleaved text wrongly shares index 1 with tool call 1. + frames := []*schema.AgenticMessage{ + streamBlock(0, schema.NewContentBlock(&schema.AssistantGenText{Text: "\n\n"})), + streamBlock(1, schema.NewContentBlock(&schema.FunctionToolCall{CallID: "call_1", Name: "exec", Arguments: ""})), + streamBlock(1, schema.NewContentBlock(&schema.FunctionToolCall{Arguments: `{"command": "whoami"}`})), + streamBlock(1, schema.NewContentBlock(&schema.AssistantGenText{Text: "\n"})), + streamBlock(2, schema.NewContentBlock(&schema.FunctionToolCall{CallID: "call_2", Name: "exec", Arguments: ""})), + streamBlock(2, schema.NewContentBlock(&schema.FunctionToolCall{Arguments: `{"command": "id"}`})), + } + + if _, err := schema.ConcatAgenticMessages(frames); err == nil { + t.Fatal("precondition failed: broken frames should fail ConcatAgenticMessages") + } + + repaired := newAgenticStreamBlockIndexRepairModel(&agenticBlockIndexFakeModel{frames: frames}) + sr, err := repaired.Stream(context.Background(), nil) + if err != nil { + t.Fatalf("Stream: %v", err) + } + var got []*schema.AgenticMessage + for { + frame, rerr := sr.Recv() + if rerr == io.EOF { + break + } + if rerr != nil { + t.Fatalf("Recv: %v", rerr) + } + got = append(got, frame) + } + + msg, err := schema.ConcatAgenticMessages(got) + if err != nil { + t.Fatalf("ConcatAgenticMessages after repair: %v", err) + } + if msg == nil { + t.Fatal("aggregated message is nil after repair") + } + var toolCalls int + var text string + for _, b := range msg.ContentBlocks { + if b.FunctionToolCall != nil { + toolCalls++ + } + if b.AssistantGenText != nil { + text += b.AssistantGenText.Text + } + } + if toolCalls != 2 { + t.Fatalf("expected 2 tool calls, got %d", toolCalls) + } + if text != "\n\n\n" { + t.Fatalf("expected concatenated text %q, got %q", "\n\n\n", text) + } +} + +func TestAgenticStreamBlockIndexRepair_ValidStreamUntouched(t *testing.T) { + frames := []*schema.AgenticMessage{ + streamBlock(0, schema.NewContentBlock(&schema.Reasoning{Text: "思考"})), + streamBlock(1, schema.NewContentBlock(&schema.AssistantGenText{Text: "你好"})), + streamBlock(1, schema.NewContentBlock(&schema.AssistantGenText{Text: "世界"})), + } + repaired := newAgenticStreamBlockIndexRepairModel(&agenticBlockIndexFakeModel{frames: frames}) + sr, err := repaired.Stream(context.Background(), nil) + if err != nil { + t.Fatalf("Stream: %v", err) + } + var got []*schema.AgenticMessage + for { + frame, rerr := sr.Recv() + if rerr == io.EOF { + break + } + if rerr != nil { + t.Fatalf("Recv: %v", rerr) + } + got = append(got, frame) + } + for i, frame := range got { + for j, b := range frame.ContentBlocks { + if b.StreamingMeta == nil || frames[i].ContentBlocks[j].StreamingMeta == nil { + continue + } + if b.StreamingMeta.Index != frames[i].ContentBlocks[j].StreamingMeta.Index { + t.Fatalf("frame %d block %d index changed: %d -> %d", i, j, + frames[i].ContentBlocks[j].StreamingMeta.Index, b.StreamingMeta.Index) + } + } + } +} diff --git a/internal/multiagent/eino_model_resilience.go b/internal/multiagent/eino_model_resilience.go index 4725a3390..c1a0064c2 100644 --- a/internal/multiagent/eino_model_resilience.go +++ b/internal/multiagent/eino_model_resilience.go @@ -102,7 +102,11 @@ func newEinoAgenticChatModelFactory( return nil, fmt.Errorf("eino agentic model: provider %q is not supported", strings.TrimSpace(oa.Provider)) } if isEinoAgenticClaudeProvider(oa.Provider) { - return newEinoClaudeAgenticChatModel(ctx, oa, mode, baseHTTPClient, reasoningClient) + nativeModel, err := newEinoClaudeAgenticChatModel(ctx, oa, mode, baseHTTPClient, reasoningClient) + if err != nil { + return nil, err + } + return newAgenticStreamBlockIndexRepairModel(nativeModel), nil } httpClient := openai.NewEinoHTTPClient(&oa, baseHTTPClient) openai.AttachSummarizationDiagTransport(httpClient, logger) @@ -118,7 +122,11 @@ func newEinoAgenticChatModelFactory( if mode == einoModelModePlanner { modelCfg.ExtraFields = reasoning.AgenticOpenAIPlannerExtraFields(&oa) } - return agenticopenai.NewChatModel(ctx, modelCfg) + chatModel, err := agenticopenai.NewChatModel(ctx, modelCfg) + if err != nil { + return nil, err + } + return newAgenticStreamBlockIndexRepairModel(chatModel), nil } } diff --git a/internal/multiagent/eino_model_resilience_test.go b/internal/multiagent/eino_model_resilience_test.go index 8f7f51089..b5950da87 100644 --- a/internal/multiagent/eino_model_resilience_test.go +++ b/internal/multiagent/eino_model_resilience_test.go @@ -8,7 +8,6 @@ import ( "cyberstrike-ai/internal/config" - agenticclaude "github.com/cloudwego/eino-ext/components/model/agenticclaude" "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" @@ -354,8 +353,8 @@ func TestNewEinoAgenticChatModelFactoryBuildsNativeClaudeBackend(t *testing.T) { if m == nil { t.Fatal("claude agentic factory returned nil model") } - if _, ok := m.(*agenticclaude.Model); !ok { - t.Fatalf("claude agentic factory returned %T, want native agenticclaude.Model", m) + if _, ok := m.(*agenticStreamBlockIndexRepairModel); !ok { + t.Fatalf("claude agentic factory returned %T, want block-index-repair wrapper around native agenticclaude.Model", m) } gate := evaluateEinoAgenticModelGate(agenticModelGateFactory(factory, config.OpenAIConfig{ Provider: "claude",