diff --git a/internal/multiagent/context_budget.go b/internal/multiagent/context_budget.go new file mode 100644 index 00000000..ba93e3a9 --- /dev/null +++ b/internal/multiagent/context_budget.go @@ -0,0 +1,328 @@ +package multiagent + +import ( + "context" + "fmt" + "strings" + "unicode/utf8" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/middlewares/summarization" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +const ( + toolOutputTruncationMarker = "\n\n...[tool output truncated; full text persisted in reduction cache or summarization transcript]...\n\n" + aggressiveToolTruncDivisor = 4 +) + +// isEinoContextOverflowError reports API-side context window rejections. +func isEinoContextOverflowError(err error) bool { + if err == nil { + return false + } + msg := strings.ToLower(strings.TrimSpace(err.Error())) + if msg == "" { + return false + } + markers := []string{ + "context length", + "context_length", + "maximum context", + "max context", + "context window", + "context overflow", + "too many tokens", + "token limit", + "tokens exceed", + "exceeds the context", + "input is too long", + "prompt is too long", + "request too large", + } + for _, m := range markers { + if strings.Contains(msg, m) { + return true + } + } + return false +} + +func truncateBytesWithMarker(content string, maxBytes int, marker string) string { + if maxBytes <= 0 || len(content) <= maxBytes { + return content + } + if marker == "" { + marker = toolOutputTruncationMarker + } + budget := maxBytes - len(marker) + if budget <= 0 { + if len(marker) > maxBytes { + return marker[:maxBytes] + } + return marker + } + head := budget / 2 + tail := budget - head + for head > 0 && !utf8.RuneStart(content[head]) { + head-- + } + tailStart := len(content) - tail + for tailStart < len(content) && !utf8.RuneStart(content[tailStart]) { + tailStart++ + } + return content[:head] + marker + content[tailStart:] +} + +func cloneMessage(msg adk.Message) adk.Message { + if msg == nil { + return nil + } + cloned := *msg + return &cloned +} + +func truncateMessageToolContent(msg adk.Message, maxBytes int, spillRef string) adk.Message { + if msg == nil || maxBytes <= 0 { + return msg + } + out := cloneMessage(msg) + marker := toolOutputTruncationMarker + if spillRef != "" { + marker = fmt.Sprintf("\n\n...[tool output truncated; retrieve full text via: %s]...\n\n", spillRef) + } + switch out.Role { + case schema.Tool: + out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker) + case schema.Assistant: + if out.ReasoningContent != "" { + out.ReasoningContent = truncateBytesWithMarker(out.ReasoningContent, maxBytes, marker) + } + if out.Content != "" { + out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker) + } + case schema.User: + if out.Content != "" { + out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker) + } + } + return out +} + +func countMessagesTokens( + ctx context.Context, + msgs []adk.Message, + counter summarization.TokenCounterFunc, + tools []*schema.ToolInfo, +) (int, error) { + if counter == nil { + return 0, nil + } + n, err := counter(ctx, &summarization.TokenCounterInput{Messages: msgs, Tools: tools}) + if err != nil { + return 0, err + } + return n, nil +} + +func truncateRoundMessagesToTokenBudget( + ctx context.Context, + round messageRound, + tokenBudget int, + counter summarization.TokenCounterFunc, + toolMaxBytes int, + spillRef string, +) ([]adk.Message, error) { + if tokenBudget <= 0 || len(round.messages) == 0 { + return nil, nil + } + msgs := append([]adk.Message(nil), round.messages...) + if n, err := countMessagesTokens(ctx, msgs, counter, nil); err != nil { + return nil, err + } else if n <= tokenBudget { + return msgs, nil + } + if toolMaxBytes <= 0 { + toolMaxBytes = 12000 + } + for pass := 0; pass < 8 && toolMaxBytes >= 32; pass++ { + out := make([]adk.Message, 0, len(msgs)) + for _, msg := range msgs { + switch { + case msg != nil && msg.Role == schema.Tool: + out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef)) + case msg != nil && msg.Role == schema.Assistant: + out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef)) + default: + out = append(out, msg) + } + } + n, err := countMessagesTokens(ctx, out, counter, nil) + if err != nil { + return nil, err + } + if n <= tokenBudget { + return out, nil + } + msgs = out + toolMaxBytes /= 2 + } + return msgs, nil +} + +type compactMessagesOpts struct { + maxTokens int + counter summarization.TokenCounterFunc + toolMaxBytes int + spillRef string + aggressive bool + logger *zap.Logger + phase string +} + +func compactMessagesByDroppingRounds( + ctx context.Context, + messages []adk.Message, + opts compactMessagesOpts, +) ([]adk.Message, bool) { + if opts.maxTokens <= 0 || len(messages) == 0 || opts.counter == nil { + return messages, false + } + before, err := countMessagesTokens(ctx, messages, opts.counter, nil) + if err != nil || before <= opts.maxTokens { + return messages, false + } + + systems := make([]adk.Message, 0, 1) + contextMsgs := make([]adk.Message, 0, len(messages)) + for _, msg := range messages { + if msg != nil && msg.Role == schema.System && len(contextMsgs) == 0 { + systems = append(systems, msg) + continue + } + if msg != nil { + contextMsgs = append(contextMsgs, msg) + } + } + rounds := splitMessagesIntoRounds(contextMsgs) + if len(rounds) == 0 { + return messages, false + } + + startIdx := 0 + if opts.aggressive { + startIdx = len(rounds) - 1 + if startIdx < 0 { + startIdx = 0 + } + } + dropped := 0 + for len(rounds) > 1 || (opts.aggressive && len(rounds) == 1) { + if !opts.aggressive && len(rounds) <= 1 { + break + } + if opts.aggressive && len(rounds) == 1 { + // Fall through to latest-round truncation below. + break + } + rounds = rounds[1:] + dropped++ + candidate := append([]adk.Message(nil), systems...) + for _, round := range rounds { + candidate = append(candidate, round.messages...) + } + after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil) + if countErr != nil { + break + } + if after <= opts.maxTokens { + if opts.logger != nil { + opts.logger.Warn("eino context compacted by dropping older rounds", + zap.String("phase", opts.phase), + zap.Int("tokens_before", before), + zap.Int("tokens_after", after), + zap.Int("max_tokens", opts.maxTokens), + zap.Int("dropped_rounds", dropped), + zap.Bool("aggressive", opts.aggressive), + ) + } + return candidate, true + } + if opts.aggressive { + break + } + } + + if len(rounds) == 0 { + return messages, false + } + latest := rounds[len(rounds)-1] + truncated, truncErr := truncateRoundMessagesToTokenBudget( + ctx, latest, opts.maxTokens, opts.counter, opts.toolMaxBytes, opts.spillRef, + ) + if truncErr != nil || len(truncated) == 0 { + if opts.logger != nil { + opts.logger.Warn("eino context still above budget after round compaction; passing through without local error", + zap.String("phase", opts.phase), + zap.Int("tokens_before", before), + zap.Int("max_tokens", opts.maxTokens), + zap.Bool("aggressive", opts.aggressive), + ) + } + return messages, false + } + candidate := append([]adk.Message(nil), systems...) + if dropped > 0 || startIdx > 0 { + for _, round := range rounds[:len(rounds)-1] { + candidate = append(candidate, round.messages...) + } + } + candidate = append(candidate, truncated...) + after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil) + if countErr != nil { + return messages, false + } + if opts.logger != nil { + opts.logger.Warn("eino context compacted by truncating latest round tool output", + zap.String("phase", opts.phase), + zap.Int("tokens_before", before), + zap.Int("tokens_after", after), + zap.Int("max_tokens", opts.maxTokens), + zap.Int("dropped_rounds", dropped), + zap.Bool("aggressive", opts.aggressive), + ) + } + return candidate, true +} + +func aggressiveCompactMessagesForOverflow( + ctx context.Context, + messages []adk.Message, + maxTotalTokens int, + modelName string, + toolMaxBytes int, + phase string, + logger *zap.Logger, +) []adk.Message { + if len(messages) == 0 || maxTotalTokens <= 0 { + return messages + } + budget := maxTotalTokens * 70 / 100 + if budget < 4096 { + budget = 4096 + } + aggressiveToolMax := toolMaxBytes / aggressiveToolTruncDivisor + if aggressiveToolMax < 2048 { + aggressiveToolMax = 2048 + } + out, _ := compactMessagesByDroppingRounds(ctx, messages, compactMessagesOpts{ + maxTokens: budget, + counter: einoSummarizationTokenCounter(modelName), + toolMaxBytes: aggressiveToolMax, + aggressive: true, + logger: logger, + phase: phase, + }) + return out +} diff --git a/internal/multiagent/context_budget_test.go b/internal/multiagent/context_budget_test.go new file mode 100644 index 00000000..a9b0ece8 --- /dev/null +++ b/internal/multiagent/context_budget_test.go @@ -0,0 +1,104 @@ +package multiagent + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestIsEinoContextOverflowError(t *testing.T) { + t.Parallel() + cases := []struct { + err error + want bool + }{ + {nil, false}, + {errors.New("context length exceeded"), true}, + {errors.New("maximum context length"), true}, + {errors.New("input is too long for model"), true}, + {errors.New("HTTP 429 Too Many Requests"), false}, + {errors.New("invalid api key"), false}, + } + for _, tc := range cases { + if got := isEinoContextOverflowError(tc.err); got != tc.want { + t.Fatalf("isEinoContextOverflowError(%v) = %v, want %v", tc.err, got, tc.want) + } + } +} + +func TestTruncateRoundMessagesToTokenBudget(t *testing.T) { + huge := strings.Repeat("x", 8000) + round := messageRound{messages: []adk.Message{ + assistantToolCallsMsg("", "c1"), + schema.ToolMessage(huge, "c1"), + }} + out, err := truncateRoundMessagesToTokenBudget( + context.Background(), round, 256, einoSummarizationTokenCounter("gpt-4o"), 512, "", + ) + if err != nil { + t.Fatal(err) + } + for _, msg := range out { + if msg != nil && msg.Role == schema.Tool && len(msg.Content) >= len(huge) { + t.Fatalf("expected truncated tool output, got len=%d", len(msg.Content)) + } + } +} + +func TestBuildBudgetedSummarizationModelInputTruncatesOversizedLatestRound(t *testing.T) { + huge := strings.Repeat("x", 8000) + msgs := []adk.Message{ + assistantToolCallsMsg("", "call-latest"), + schema.ToolMessage(huge, "call-latest"), + } + counter := einoSummarizationTokenCounter("gpt-4o") + input, dropped, err := buildBudgetedSummarizationModelInput( + context.Background(), + schema.SystemMessage("sys"), + schema.UserMessage("instr"), + msgs, + counter, + 512, + summarizationInputBudgetOpts{toolMaxBytes: 256}, + ) + if err != nil { + t.Fatal(err) + } + if dropped != 0 { + t.Fatalf("expected no dropped rounds, got %d", dropped) + } + toolContent := "" + for _, msg := range input { + if msg != nil && msg.Role == schema.Tool { + toolContent = msg.Content + } + } + if len(toolContent) >= len(huge) { + t.Fatalf("expected oversized tool output to be compacted, got len=%d", len(toolContent)) + } +} + +func TestModelInputSoftBudgetNeverErrors(t *testing.T) { + mw := &modelInputSoftBudgetMiddleware{ + maxTokens: 4, + toolMaxBytes: 16, + counter: fixedTokenCounter(4), + phase: "test", + } + state := &adk.ChatModelAgentState{Messages: []adk.Message{ + schema.UserMessage("u"), + assistantToolCallsMsg("", "c1"), + schema.ToolMessage(strings.Repeat("t", 200), "c1"), + }} + _, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatalf("soft budget must not error: %v", err) + } + if out == nil || len(out.Messages) == 0 { + t.Fatal("expected compacted messages") + } +} diff --git a/internal/multiagent/continuation_user_dedup_middleware.go b/internal/multiagent/continuation_user_dedup_middleware.go new file mode 100644 index 00000000..fdb3b915 --- /dev/null +++ b/internal/multiagent/continuation_user_dedup_middleware.go @@ -0,0 +1,104 @@ +package multiagent + +import ( + "context" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// continuationSessionMarker matches Cursor / IDE session-resume user injections. +const continuationSessionMarker = "This session is being continued from a previous conversation" + +// continuationUserDedupMiddleware keeps only the latest session-resume user message when +// multiple continuation injections were stacked (e.g. after repeated out-of-context resumes). +type continuationUserDedupMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + phase string +} + +func newContinuationUserDedupMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware { + return &continuationUserDedupMiddleware{logger: logger, phase: phase} +} + +func (m *continuationUserDedupMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.ChatModelAgentState, + mc *adk.ModelContext, +) (context.Context, *adk.ChatModelAgentState, error) { + _ = mc + if m == nil || state == nil || len(state.Messages) == 0 { + return ctx, state, nil + } + deduped, dropped := dedupContinuationUserMessages(state.Messages) + if dropped == 0 { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Info("eino continuation user messages deduplicated", + zap.String("phase", m.phase), + zap.Int("dropped", dropped), + zap.Int("messages_before", len(state.Messages)), + zap.Int("messages_after", len(deduped)), + ) + } + out := *state + out.Messages = deduped + return ctx, &out, nil +} + +func adkUserMessageText(msg adk.Message) string { + if msg == nil { + return "" + } + var b strings.Builder + if s := strings.TrimSpace(msg.Content); s != "" { + b.WriteString(s) + } + for _, part := range msg.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText { + if s := strings.TrimSpace(part.Text); s != "" { + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(s) + } + } + } + return b.String() +} + +func isContinuationUserMessage(msg adk.Message) bool { + if msg == nil || msg.Role != schema.User { + return false + } + return strings.Contains(adkUserMessageText(msg), continuationSessionMarker) +} + +func dedupContinuationUserMessages(msgs []adk.Message) ([]adk.Message, int) { + lastIdx := -1 + contCount := 0 + for i, msg := range msgs { + if !isContinuationUserMessage(msg) { + continue + } + contCount++ + lastIdx = i + } + if contCount <= 1 { + return msgs, 0 + } + out := make([]adk.Message, 0, len(msgs)-(contCount-1)) + dropped := 0 + for i, msg := range msgs { + if isContinuationUserMessage(msg) && i != lastIdx { + dropped++ + continue + } + out = append(out, msg) + } + return out, dropped +} diff --git a/internal/multiagent/continuation_user_dedup_middleware_test.go b/internal/multiagent/continuation_user_dedup_middleware_test.go new file mode 100644 index 00000000..75987d86 --- /dev/null +++ b/internal/multiagent/continuation_user_dedup_middleware_test.go @@ -0,0 +1,65 @@ +package multiagent + +import ( + "context" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func continuationUser(text string) adk.Message { + return &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: continuationSessionMarker + "\n" + text}, + {Type: schema.ChatMessagePartTypeText, Text: "Please continue the conversation from where we left it off."}, + }, + } +} + +func TestDedupContinuationUserMessages_KeepsLatest(t *testing.T) { + msgs := []adk.Message{ + continuationUser("summary old"), + schema.UserMessage("real task"), + continuationUser("summary new"), + } + out, dropped := dedupContinuationUserMessages(msgs) + if dropped != 1 { + t.Fatalf("dropped=%d want 1", dropped) + } + if len(out) != 2 { + t.Fatalf("len=%d want 2", len(out)) + } + if out[0].Role != schema.User || adkUserMessageText(out[0]) != "real task" { + t.Fatalf("first should remain real task, got %q", adkUserMessageText(out[0])) + } + if !strings.Contains(adkUserMessageText(out[1]), "summary new") { + t.Fatalf("latest continuation not kept: %q", adkUserMessageText(out[1])) + } +} + +func TestDedupContinuationUserMessages_NoOpSingle(t *testing.T) { + msgs := []adk.Message{continuationUser("only"), schema.UserMessage("task")} + out, dropped := dedupContinuationUserMessages(msgs) + if dropped != 0 || len(out) != 2 { + t.Fatalf("unexpected change dropped=%d len=%d", dropped, len(out)) + } +} + +func TestContinuationUserDedupMiddleware(t *testing.T) { + mw := newContinuationUserDedupMiddleware(nil, "test") + state := &adk.ChatModelAgentState{Messages: []adk.Message{ + continuationUser("old"), + continuationUser("new"), + schema.UserMessage("task"), + }} + _, out, err := mw.(*continuationUserDedupMiddleware).BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatal(err) + } + if len(out.Messages) != 2 { + t.Fatalf("want 2 messages after dedup, got %d", len(out.Messages)) + } +} diff --git a/internal/multiagent/eino_adk_run_loop.go b/internal/multiagent/eino_adk_run_loop.go new file mode 100644 index 00000000..2b31eebc --- /dev/null +++ b/internal/multiagent/eino_adk_run_loop.go @@ -0,0 +1,478 @@ +package multiagent + +import ( + "context" + "errors" + "fmt" + "regexp" + "strings" + "sync" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/einoobserve" + "cyberstrike-ai/internal/security" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// normalizeStreamingDelta 将可能是“累计片段”的 chunk 归一化为“纯增量”。 +// 一些模型/桥接层在流式过程中会重复发送已输出前缀,前端若直接 buffer+=chunk 会出现重复文本。 +// +// 注意:与 internal/openai.normalizeStreamingDelta 保持一致。 +func normalizeStreamingDelta(current, incoming string) (next, delta string) { + if incoming == "" { + return current, "" + } + if current == "" { + return incoming, incoming + } + if strings.HasPrefix(incoming, current) && len(incoming) > len(current) { + return incoming, incoming[len(current):] + } + if incoming == current && utf8.RuneCountInString(current) > 1 { + return current, "" + } + return current + incoming, incoming +} + +func isInterruptContinue(ctx context.Context) bool { + if ctx == nil { + return false + } + return errors.Is(context.Cause(ctx), ErrInterruptContinue) +} + +func isEinoStreamCanceled(err error) bool { + if err == nil { + return false + } + if errors.Is(err, adk.ErrStreamCanceled) { + return true + } + var streamCanceled *adk.StreamCanceledError + return errors.As(err, &streamCanceled) +} + +func isEinoCancelError(err error) bool { + if err == nil { + return false + } + var cancelErr *adk.CancelError + return errors.As(err, &cancelErr) +} + +// isEinoVoluntaryCancelErr reports cancel signals produced by Agent Cancel / +// TurnLoop preempt (CancelError, ErrStreamCanceled, context.Canceled). +func isEinoVoluntaryCancelErr(err error) bool { + if err == nil { + return false + } + return isEinoCancelError(err) || isEinoStreamCanceled(err) || errors.Is(err, context.Canceled) +} + +// isEinoTurnLoopPreemptErr is true when a cancel/stream-cancel leaked from the +// current agent turn while the host task context is still alive. TurnLoop +// interrupt-continue does not cancel the parent context; treating that leak as +// fatal would abort the whole run instead of starting the queued next turn. +func isEinoTurnLoopPreemptErr(ctx context.Context, err error) bool { + if err == nil || !isEinoVoluntaryCancelErr(err) { + return false + } + if ctx != nil && ctx.Err() != nil { + return false + } + return true +} + +func isEinoIterationLimitError(err error) bool { + if err == nil { + return false + } + msg := strings.ToLower(strings.TrimSpace(err.Error())) + if msg == "" { + return false + } + return strings.Contains(msg, "max iteration") || + strings.Contains(msg, "maximum iteration") || + strings.Contains(msg, "maximum iterations") || + strings.Contains(msg, "iteration limit") || + strings.Contains(msg, "达到最大迭代") +} + +// einoADKRunLoopArgs 将 Eino adk.Runner 事件循环从 RunDeepAgent / RunEinoSingleChatModelAgent 中抽出复用。 +type einoADKRunLoopArgs struct { + OrchMode string + OrchestratorName string + ConversationID string + Progress func(eventType, message string, data interface{}) + Logger *zap.Logger + SnapshotMCPIDs func() []string + StreamsMainAssistant func(agent string) bool + EinoRoleTag func(agent string) string + CheckpointDir string + // RunRetryMaxAttempts / RunRetryMaxBackoffSec:429、5xx、网络抖动时的指数退避续跑(0=默认 4 次 / 30s 上限)。 + RunRetryMaxAttempts int + RunRetryMaxBackoffSec int + + McpIDsMu *sync.Mutex + McpIDs *[]string + + // FilesystemMonitorAgent / FilesystemMonitorRecord 非 nil 时,将 Eino ADK filesystem 中间件工具(ls/read_file/write_file/edit_file/glob/grep) + // 在完成时写入 MCP 监控;execute 仍由 eino_execute_monitor 记录,此处跳过。 + FilesystemMonitorAgent *agent.Agent + FilesystemMonitorRecord einomcp.ExecutionRecorder + MCPExecutionBinder *MCPExecutionBinder + + // ToolInvokeNotify 与 einomcp.ToolsFromDefinitions 共享:run loop 在迭代前 Set,execute/MCP 桥 Fire 时立即推送 tool_result(ADK 晚到经 toolResultEmitter 去重)。 + ToolInvokeNotify *einomcp.ToolInvokeNotifyHolder + + DA adk.Agent + + // EmptyResponseMessage 当未捕获到助手正文时的占位(多代理与单代理文案不同)。 + EmptyResponseMessage string + + // ModelFacingTrace 可选:由各 ChatModelAgent Handlers 链末尾中间件写入「即将送入模型」的消息快照; + // 非空时优先用于 LastAgentTraceInput 序列化,使续跑与 summarization/reduction 后的上下文一致。 + ModelFacingTrace *modelFacingTraceHolder + + // EinoCallbacks 可选:为 ADK Runner 注入 eino [callbacks] 全链路观测(见 internal/einoobserve)。 + EinoCallbacks *config.MultiAgentEinoCallbacksConfig + + // MaxTotalTokens / ToolMaxBytes / ModelName 用于 context overflow 时的激进压缩续跑。 + MaxTotalTokens int + ToolMaxBytes int + ModelName string + MiddlewareConfig *config.MultiAgentEinoMiddlewareConfig + + // TurnLoopInterruptTimeout 仅供测试/特殊运行时覆盖;0 使用 EinoTurnLoopRuntime 默认值。 + TurnLoopInterruptTimeout time.Duration +} + +func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs []adk.Message) (*RunResult, error) { + if args == nil || args.DA == nil { + return nil, fmt.Errorf("eino run loop: args 或 Agent 为空") + } + if args.McpIDs == nil { + s := []string{} + args.McpIDs = &s + } + if args.McpIDsMu == nil { + args.McpIDsMu = &sync.Mutex{} + } + + orchMode := args.OrchMode + orchestratorName := args.OrchestratorName + conversationID := args.ConversationID + progress := args.Progress + logger := args.Logger + runID := newEinoRunID() + progress = withEinoRunIDProgress(runID, progress) + args.Progress = progress + if logger != nil { + logger.Info("eino run session started", + zap.String("runId", runID), + zap.String("conversationId", conversationID), + zap.String("orchestration", orchMode), + zap.String("orchestratorName", orchestratorName), + ) + } + snapshotMCPIDs := args.SnapshotMCPIDs + if snapshotMCPIDs == nil { + snapshotMCPIDs = func() []string { return nil } + } + streamsMainAssistant := args.StreamsMainAssistant + if streamsMainAssistant == nil { + streamsMainAssistant = func(agent string) bool { + return agent == "" || agent == orchestratorName + } + } + einoRoleTag := args.EinoRoleTag + if einoRoleTag == nil { + einoRoleTag = func(agent string) string { + if streamsMainAssistant(agent) { + return "orchestrator" + } + return "sub" + } + } + // panic recovery:防止 Eino 框架内部 panic 导致整个 goroutine 崩溃、连接无法正常关闭。 + defer func() { + if r := recover(); r != nil { + if logger != nil { + logger.Error("eino runner panic recovered", zap.Any("recover", r), zap.Stack("stack")) + } + if progress != nil { + progress("error", fmt.Sprintf("Internal error: %v / 内部错误: %v", r, r), map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + }) + } + } + }() + + msgs := append([]adk.Message(nil), baseMsgs...) + + emptyHint := strings.TrimSpace(args.EmptyResponseMessage) + if emptyHint == "" { + emptyHint = "(Eino session completed but no assistant text was captured. Check process details or logs.) " + + "(Eino 会话已完成,但未捕获到助手文本输出。请查看过程详情或日志。)" + } + + if args.EinoCallbacks != nil { + ctx = einoobserve.AttachAgentRunCallbacks(ctx, args.EinoCallbacks, einoobserve.Params{ + Logger: logger, + Progress: progress, + ConversationID: conversationID, + OrchMode: orchMode, + OrchestratorName: orchestratorName, + RunID: runID, + }) + } + + drain := newEinoRunEventDrain(einoRunEventDrainConfig{ + Context: ctx, + ConversationID: conversationID, + OrchMode: orchMode, + OrchestratorName: orchestratorName, + Progress: progress, + Logger: logger, + BaseMessages: msgs, + SnapshotMCPIDs: snapshotMCPIDs, + StreamsMainAssistant: streamsMainAssistant, + EinoRoleTag: einoRoleTag, + MiddlewareConfig: args.MiddlewareConfig, + FilesystemMonitorAgent: args.FilesystemMonitorAgent, + FilesystemMonitorRecord: args.FilesystemMonitorRecord, + MCPExecutionBinder: args.MCPExecutionBinder, + }) + session := newEinoRunRuntimeSession(einoRunRuntimeSessionConfig{ + Context: ctx, + Args: args, + Drain: drain, + BaseMessages: msgs, + EmptyHint: emptyHint, + SnapshotMCPIDs: snapshotMCPIDs, + EinoRoleTag: einoRoleTag, + }) + defer session.Close() + + // 仅在退避重试后真正收到数据/完成一步时清零,避免重启后首个无错 ADK 事件误把计数打回 0。 + drain.BindHandlers(session.ConfirmRecovery) + + for { + // iter.Next 可能长时间阻塞(工具执行、模型推理);须与 ctx 联动,否则取消/超时无法及时 flush pending。 + ev, ok, iterCtxErr := nextAgentEventWithContext(ctx, session.Iterator()) + if iterCtxErr != nil { + return session.HandleIteratorContextError(iterCtxErr) + } + if !ok { + // iter 结束并不总是“正常完成”: + // 当取消/超时发生在 iter.Next() 阻塞期间时,可能直接返回 !ok。 + // 此时必须保留 checkpoint,避免后续恢复时被误判为“无断点”而全量重跑。 + completed, result, err := session.HandleIteratorEnd() + if result != nil || err != nil { + return result, err + } + if completed { + break + } + continue + } + if ev == nil { + continue + } + if ev.Err != nil { + handled := session.HandleRunError(ev.Err) + if handled.Result != nil || handled.Err != nil { + return handled.Result, handled.Err + } + if handled.Restarted { + continue + } + } + drain.ObserveAgent(ev.AgentName) + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + mv := ev.Output.MessageOutput + + if drain.HandleToolResultStreaming(mv, ev.AgentName) { + continue + } + + if handledStream, streamRecvErr := drain.HandleAssistantStream(mv, ev.AgentName); handledStream { + if streamRecvErr != nil { + handled := session.HandleStreamError(streamRecvErr, ev.AgentName) + if handled.Result != nil || handled.Err != nil { + return handled.Result, handled.Err + } + if handled.Restarted { + continue + } + } else { + session.ConfirmRecovery() + } + continue + } + + msg, gerr := mv.GetMessage() + if gerr != nil || msg == nil { + continue + } + drain.HandleMaterialized(mv, msg, ev.AgentName) + session.ConfirmRecovery() + } + + return session.BuildFinalResult(), nil +} + +// modelFacingTraceSnapshot returns only the state that actually reached the model boundary. +// Never fall back to event-stream accumulation here: it can contain pre-reduction tool output +// that the model never received (for example when summarization failed before the first call). +func modelFacingTraceSnapshot(args *einoADKRunLoopArgs) []adk.Message { + if args != nil && args.ModelFacingTrace != nil { + if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 { + return snap + } + } + return nil +} + +// friendlyEinoExecuteInvokeTail 将 Eino execute 超时/中断/流异常转为简短提示。 +// 命令非零退出(ExecuteExitError)已有 exec 对齐的正文,不再追加「执行未正常结束」。 +func friendlyEinoExecuteInvokeTail(invokeErr error) string { + if invokeErr == nil { + return "" + } + var exitErr *ExecuteExitError + if errors.As(invokeErr, &exitErr) { + return "" + } + if errors.Is(invokeErr, context.DeadlineExceeded) { + return einoExecuteTimeoutUserHint() + } + if errors.Is(invokeErr, context.Canceled) { + return "" + } + if strings.Contains(invokeErr.Error(), "shell inactivity timeout") { + return "" + } + return "[执行未正常结束] " + invokeErr.Error() +} + +// einoToolResultIsError 统一判断 Eino 工具结果是否应标记为错误(与 MCP exec 的 IsError 对齐)。 +func einoToolResultIsError(toolName, content string) bool { + if strings.HasPrefix(content, einomcp.ToolErrorPrefix) { + return true + } + if strings.TrimSpace(toolName) == "execute" && security.IsCommandFailureResult(content) { + return true + } + return false +} + +func isMCPBackgroundWaitResult(content string) bool { + text := strings.ToLower(strings.TrimSpace(content)) + if text == "" { + return false + } + hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`) + hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") || + strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) || + strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`) + hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") || + strings.Contains(text, "本次等待已到达") || + strings.Contains(text, "wait_timeout:") || + strings.Contains(text, "background execution") || + strings.Contains(text, "still running") || + strings.Contains(text, "仍未完成") + return hasExecutionID && hasRunningStatus && hasSoftWaitSignal +} + +func mcpExecutionIDFromWaitResult(content string) string { + re := regexp.MustCompile(`(?i)"?execution_id"?\s*[:=]\s*"?([0-9a-f]{8}-[0-9a-f-]{12,})"?`) + if m := re.FindStringSubmatch(content); len(m) > 1 { + return strings.TrimSpace(m[1]) + } + for _, line := range strings.Split(content, "\n") { + line = strings.TrimSpace(line) + lower := strings.ToLower(line) + if !strings.HasPrefix(lower, "execution_id:") { + continue + } + return strings.Trim(strings.TrimSpace(line[len("execution_id:"):]), `"'`) + } + return "" +} + +// einoToolResultBody 去掉工具错误前缀,返回展示/持久化正文。 +func einoToolResultBody(content string) string { + if strings.HasPrefix(content, einomcp.ToolErrorPrefix) { + return strings.TrimPrefix(content, einomcp.ToolErrorPrefix) + } + return content +} + +// nextAgentEventWithContext 在 ctx 取消时不再无限阻塞于 iter.Next()(工具执行/模型推理期间常见)。 +func nextAgentEventWithContext(ctx context.Context, iter *adk.AsyncIterator[*adk.AgentEvent]) (ev *adk.AgentEvent, ok bool, ctxErr error) { + if iter == nil { + return nil, false, nil + } + type nextRes struct { + ev *adk.AgentEvent + ok bool + } + ch := make(chan nextRes, 1) + go func() { + e, o := iter.Next() + ch <- nextRes{e, o} + }() + select { + case <-ctx.Done(): + return nil, false, ctx.Err() + case res := <-ch: + return res.ev, res.ok, nil + } +} + +// recvSchemaMessageStream 消费 ADK Tool 流式结果;ctx 取消时立即返回,避免 amass 等无输出时永久阻塞。 +func recvSchemaMessageStream(ctx context.Context, stream *schema.StreamReader[*schema.Message]) (content, toolCallID, toolName string, recvErr error) { + if stream == nil { + return "", "", "", nil + } + var buf strings.Builder + recvErr = recvEinoSchemaMessageStreamWithContext(ctx, stream, 8, func(chunk *schema.Message) { + if chunk.Content != "" { + buf.WriteString(chunk.Content) + } + if tid := strings.TrimSpace(chunk.ToolCallID); tid != "" { + toolCallID = tid + } + if name := strings.TrimSpace(chunk.ToolName); name != "" { + toolName = name + } + }) + return buf.String(), toolCallID, toolName, recvErr +} + +func buildEinoCheckpointID(orchMode string) string { + mode := sanitizeEinoPathSegment(strings.TrimSpace(orchMode)) + if mode == "" { + mode = "default" + } + return "runner-" + mode +} + +func buildEinoTurnLoopCheckpointID(orchMode string) string { + mode := sanitizeEinoPathSegment(strings.TrimSpace(orchMode)) + if mode == "" { + mode = "default" + } + return "turn-loop-" + mode +} diff --git a/internal/multiagent/eino_adk_run_loop_stream_test.go b/internal/multiagent/eino_adk_run_loop_stream_test.go new file mode 100644 index 00000000..2ca8381a --- /dev/null +++ b/internal/multiagent/eino_adk_run_loop_stream_test.go @@ -0,0 +1,122 @@ +package multiagent + +import ( + "context" + "errors" + "io" + "testing" + "time" + + "github.com/cloudwego/eino/schema" +) + +func TestRecvSchemaMessageStream_EOF(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + _ = sw.Send(schema.ToolMessage("hello", "tc-1"), nil) + sw.Close() + + content, tid, toolName, err := recvSchemaMessageStream(context.Background(), sr) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if content != "hello" { + t.Fatalf("content=%q want hello", content) + } + if tid != "tc-1" { + t.Fatalf("toolCallID=%q want tc-1", tid) + } + if toolName != "" { + t.Fatalf("toolName=%q want empty", toolName) + } +} + +func TestRecvSchemaMessageStream_CapturesToolName(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + _ = sw.Send(schema.ToolMessage("hello", "tc-1", schema.WithToolName("execute")), nil) + sw.Close() + + content, tid, toolName, err := recvSchemaMessageStream(context.Background(), sr) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if content != "hello" || tid != "tc-1" || toolName != "execute" { + t.Fatalf("content=%q tid=%q toolName=%q", content, tid, toolName) + } +} + +func TestRecvSchemaMessageStream_ContextCancel(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + t.Cleanup(func() { sw.Close() }) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(30 * time.Millisecond) + cancel() + }() + + content, _, _, err := recvSchemaMessageStream(ctx, sr) + if !errors.Is(err, context.Canceled) { + t.Fatalf("want context.Canceled, got %v content=%q", err, content) + } +} + +func TestRecvSchemaMessageStream_RecvError(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + want := errors.New("stream broken") + _ = sw.Send(nil, want) + sw.Close() + + _, _, _, err := recvSchemaMessageStream(context.Background(), sr) + if !errors.Is(err, want) { + t.Fatalf("want %v, got %v", want, err) + } +} + +func TestRecvSchemaMessageStream_NilStream(t *testing.T) { + content, tid, toolName, err := recvSchemaMessageStream(context.Background(), nil) + if err != nil || content != "" || tid != "" || toolName != "" { + t.Fatalf("nil stream: content=%q tid=%q toolName=%q err=%v", content, tid, toolName, err) + } +} + +func TestRecvSchemaMessageStream_EOFViaEmptyRead(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + _ = sw.Send(nil, io.EOF) + sw.Close() + + _, _, _, err := recvSchemaMessageStream(context.Background(), sr) + if err != nil { + t.Fatalf("EOF should not surface as error, got %v", err) + } +} + +func TestRecvEinoSchemaMessageStreamWithContext_SkipsNilChunks(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](4) + _ = sw.Send(nil, nil) + _ = sw.Send(schema.AssistantMessage("hello", nil), nil) + sw.Close() + + var got []string + err := recvEinoSchemaMessageStreamWithContext(context.Background(), sr, 1, func(chunk *schema.Message) { + got = append(got, chunk.Content) + }) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if len(got) != 1 || got[0] != "hello" { + t.Fatalf("chunks = %#v, want [hello]", got) + } +} + +func TestRecvEinoSchemaMessageStreamWithContext_NilStream(t *testing.T) { + called := false + err := recvEinoSchemaMessageStreamWithContext(context.Background(), nil, 0, func(*schema.Message) { + called = true + }) + if err != nil { + t.Fatalf("nil stream should not error, got %v", err) + } + if called { + t.Fatal("nil stream should not call handler") + } +} diff --git a/internal/multiagent/eino_agentic_agent_adapter.go b/internal/multiagent/eino_agentic_agent_adapter.go new file mode 100644 index 00000000..cdc7b2d9 --- /dev/null +++ b/internal/multiagent/eino_agentic_agent_adapter.go @@ -0,0 +1,81 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +type einoAgenticMessageAgentAdapter struct { + inner adk.TypedAgent[*schema.AgenticMessage] +} + +func newEinoAgenticMessageAgentAdapter(inner adk.TypedAgent[*schema.AgenticMessage]) adk.Agent { + if inner == nil { + return nil + } + return &einoAgenticMessageAgentAdapter{inner: inner} +} + +func (a *einoAgenticMessageAgentAdapter) Name(ctx context.Context) string { + if a == nil || a.inner == nil { + return "" + } + return a.inner.Name(ctx) +} + +func (a *einoAgenticMessageAgentAdapter) Description(ctx context.Context) string { + if a == nil || a.inner == nil { + return "" + } + return a.inner.Description(ctx) +} + +func (a *einoAgenticMessageAgentAdapter) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + return a.runTyped(ctx, input, nil, opts...) +} + +func (a *einoAgenticMessageAgentAdapter) Resume(ctx context.Context, info *adk.ResumeInfo, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + return a.runTyped(ctx, nil, info, opts...) +} + +func (a *einoAgenticMessageAgentAdapter) runTyped(ctx context.Context, input *adk.AgentInput, resumeInfo *adk.ResumeInfo, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer gen.Close() + if a == nil || a.inner == nil { + gen.Send(&adk.AgentEvent{Err: fmt.Errorf("agentic adapter: inner agent is nil")}) + return + } + var agenticIter *adk.AsyncIterator[*adk.TypedAgentEvent[*schema.AgenticMessage]] + if resumeInfo != nil { + resumable, ok := a.inner.(adk.TypedResumableAgent[*schema.AgenticMessage]) + if !ok { + gen.Send(&adk.AgentEvent{Err: fmt.Errorf("agentic adapter: inner agent does not support resume")}) + return + } + agenticIter = resumable.Resume(ctx, resumeInfo, opts...) + } else { + agenticInput := &adk.TypedAgentInput[*schema.AgenticMessage]{} + if input != nil { + agenticInput.EnableStreaming = input.EnableStreaming + agenticInput.Messages = EinoMessagesToAgentic(input.Messages) + } + agenticIter = a.inner.Run(ctx, agenticInput, opts...) + } + for { + ev, ok := agenticIter.Next() + if !ok { + return + } + for _, adapted := range adaptAgenticEventToEinoEvents(ev) { + if adapted != nil { + gen.Send(adapted) + } + } + } + }() + return iter +} diff --git a/internal/multiagent/eino_agentic_agent_adapter_test.go b/internal/multiagent/eino_agentic_agent_adapter_test.go new file mode 100644 index 00000000..b28aecc0 --- /dev/null +++ b/internal/multiagent/eino_agentic_agent_adapter_test.go @@ -0,0 +1,145 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +type fakeAgenticMessageAgent struct { + name string + description string + captured *adk.TypedAgentInput[*schema.AgenticMessage] + resumeInfo *adk.ResumeInfo + events []*adk.TypedAgentEvent[*schema.AgenticMessage] +} + +func (f *fakeAgenticMessageAgent) Name(context.Context) string { + return f.name +} + +func (f *fakeAgenticMessageAgent) Description(context.Context) string { + return f.description +} + +func (f *fakeAgenticMessageAgent) Run(_ context.Context, input *adk.TypedAgentInput[*schema.AgenticMessage], _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.TypedAgentEvent[*schema.AgenticMessage]] { + f.captured = input + iter, gen := adk.NewAsyncIteratorPair[*adk.TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + for _, ev := range f.events { + gen.Send(ev) + } + }() + return iter +} + +func (f *fakeAgenticMessageAgent) Resume(_ context.Context, info *adk.ResumeInfo, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.TypedAgentEvent[*schema.AgenticMessage]] { + f.resumeInfo = info + iter, gen := adk.NewAsyncIteratorPair[*adk.TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + for _, ev := range f.events { + gen.Send(ev) + } + }() + return iter +} + +func TestEinoAgenticMessageAgentAdapterConvertsInputAndEvents(t *testing.T) { + inner := &fakeAgenticMessageAgent{ + name: "agentic", + description: "typed agent", + events: []*adk.TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "agentic", + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "hello"}), + }, + }, + }, + }, + }, + }, + } + agent := newEinoAgenticMessageAgentAdapter(inner) + + if agent.Name(context.Background()) != "agentic" || agent.Description(context.Background()) != "typed agent" { + t.Fatalf("adapter metadata name=%q desc=%q", agent.Name(context.Background()), agent.Description(context.Background())) + } + iter := agent.Run(context.Background(), &adk.AgentInput{ + EnableStreaming: true, + Messages: []*schema.Message{ + schema.UserMessage("hi"), + }, + }) + + ev, ok := iter.Next() + if !ok { + t.Fatal("expected adapted event") + } + if inner.captured == nil || !inner.captured.EnableStreaming || len(inner.captured.Messages) != 1 { + t.Fatalf("captured input = %#v", inner.captured) + } + if inner.captured.Messages[0].Role != schema.AgenticRoleTypeUser || inner.captured.Messages[0].ContentBlocks[0].UserInputText.Text != "hi" { + t.Fatalf("captured message = %#v", inner.captured.Messages[0]) + } + if ev.AgentName != "agentic" || ev.Output == nil || ev.Output.MessageOutput == nil { + t.Fatalf("event = %#v", ev) + } + if ev.Output.MessageOutput.Role != schema.Assistant || ev.Output.MessageOutput.Message.Content != "hello" { + t.Fatalf("message output = %#v", ev.Output.MessageOutput) + } + if _, ok := iter.Next(); ok { + t.Fatal("expected iterator to close") + } +} + +func TestEinoAgenticMessageAgentAdapterNilInnerReturnsNil(t *testing.T) { + if got := newEinoAgenticMessageAgentAdapter(nil); got != nil { + t.Fatalf("adapter = %#v, want nil", got) + } +} + +func TestEinoAgenticMessageAgentAdapterResumeConvertsEvents(t *testing.T) { + inner := &fakeAgenticMessageAgent{ + name: "agentic", + events: []*adk.TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "agentic", + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "resumed"}), + }, + }, + }, + }, + }, + }, + } + agent, ok := newEinoAgenticMessageAgentAdapter(inner).(adk.ResumableAgent) + if !ok { + t.Fatal("adapter must implement adk.ResumableAgent") + } + info := &adk.ResumeInfo{WasInterrupted: true} + iter := agent.Resume(context.Background(), info) + ev, ok := iter.Next() + if !ok { + t.Fatal("expected adapted resume event") + } + if inner.resumeInfo != info { + t.Fatalf("resume info = %#v, want original pointer", inner.resumeInfo) + } + if ev.Output == nil || ev.Output.MessageOutput == nil || ev.Output.MessageOutput.Message.Content != "resumed" { + t.Fatalf("resume event = %#v", ev) + } +} diff --git a/internal/multiagent/eino_agentic_chat_model_agent.go b/internal/multiagent/eino_agentic_chat_model_agent.go new file mode 100644 index 00000000..fb767154 --- /dev/null +++ b/internal/multiagent/eino_agentic_chat_model_agent.go @@ -0,0 +1,64 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +type einoAgenticChatModelAgentConfig struct { + Name string + Description string + Instruction string + Model model.AgenticModel + ToolsConfig adk.ToolsConfig + MaxIterations int + Exit tool.BaseTool + + GenModelInput adk.TypedGenModelInput[*schema.AgenticMessage] + Handlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + ModelRetryConfig *adk.TypedModelRetryConfig[*schema.AgenticMessage] + ModelFailoverConfig *adk.ModelFailoverConfig[*schema.AgenticMessage] + OutputKey string +} + +func newEinoAgenticChatModelAgent(ctx context.Context, cfg einoAgenticChatModelAgentConfig) (adk.TypedResumableAgent[*schema.AgenticMessage], error) { + if cfg.Model == nil { + return nil, fmt.Errorf("eino agentic ChatModelAgent: model is required") + } + typedCfg := &adk.TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: cfg.Name, + Description: cfg.Description, + Instruction: cfg.Instruction, + Model: cfg.Model, + ToolsConfig: cfg.ToolsConfig, + MaxIterations: cfg.MaxIterations, + Exit: cfg.Exit, + GenModelInput: cfg.GenModelInput, + Handlers: cfg.Handlers, + ModelRetryConfig: cfg.ModelRetryConfig, + ModelFailoverConfig: cfg.ModelFailoverConfig, + OutputKey: cfg.OutputKey, + } + typedAgent, err := adk.NewTypedChatModelAgent(ctx, typedCfg) + if err != nil { + return nil, fmt.Errorf("eino agentic NewTypedChatModelAgent: %w", err) + } + return typedAgent, nil +} + +func newEinoAgenticChatModelAgentAdapter(ctx context.Context, cfg einoAgenticChatModelAgentConfig) (adk.Agent, error) { + typedAgent, err := newEinoAgenticChatModelAgent(ctx, cfg) + if err != nil { + return nil, err + } + agent := newEinoAgenticMessageAgentAdapter(typedAgent) + if agent == nil { + return nil, fmt.Errorf("eino agentic ChatModelAgent: adapter is nil") + } + return agent, nil +} diff --git a/internal/multiagent/eino_agentic_chat_model_agent_test.go b/internal/multiagent/eino_agentic_chat_model_agent_test.go new file mode 100644 index 00000000..63761c71 --- /dev/null +++ b/internal/multiagent/eino_agentic_chat_model_agent_test.go @@ -0,0 +1,163 @@ +package multiagent + +import ( + "context" + "strings" + "sync" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type capturingAgenticChatModel struct { + mu sync.Mutex + inputs [][]*schema.AgenticMessage + output *schema.AgenticMessage +} + +func (m *capturingAgenticChatModel) Generate(_ context.Context, input []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + m.mu.Lock() + m.inputs = append(m.inputs, input) + m.mu.Unlock() + if m.output != nil { + return m.output, nil + } + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: "agentic answer"})}, + }, nil +} + +func (m *capturingAgenticChatModel) Stream(_ context.Context, input []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + msg, err := m.Generate(context.Background(), input) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.AgenticMessage{msg}), nil +} + +func (m *capturingAgenticChatModel) snapshotInputs() [][]*schema.AgenticMessage { + m.mu.Lock() + defer m.mu.Unlock() + out := make([][]*schema.AgenticMessage, len(m.inputs)) + copy(out, m.inputs) + return out +} + +func TestNewEinoAgenticChatModelAgentAdapterRunsThroughClassicAgentBoundary(t *testing.T) { + t.Parallel() + ctx := context.Background() + trace := newModelFacingTraceHolder() + fakeModel := &capturingAgenticChatModel{} + agent, err := newEinoAgenticChatModelAgentAdapter(ctx, einoAgenticChatModelAgentConfig{ + Name: "agentic", + Description: "agentic adapter test", + Instruction: "system instruction", + Model: fakeModel, + Handlers: appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + phase: "agentic", + trace: trace, + }), + }) + if err != nil { + t.Fatalf("newEinoAgenticChatModelAgentAdapter: %v", err) + } + + iter := agent.Run(ctx, &adk.AgentInput{ + Messages: []*schema.Message{schema.UserMessage("classic input")}, + }) + var last *adk.AgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + t.Fatalf("agent event error: %v", ev.Err) + } + last = ev + } + if last == nil || last.Output == nil || last.Output.MessageOutput == nil { + t.Fatalf("last event = %#v, want message output", last) + } + if got := last.Output.MessageOutput.Message.Content; got != "agentic answer" { + t.Fatalf("classic output content = %q, want agentic answer", got) + } + + inputs := fakeModel.snapshotInputs() + if len(inputs) != 1 { + t.Fatalf("model calls = %d, want 1", len(inputs)) + } + if len(inputs[0]) != 2 { + t.Fatalf("model input messages = %d, want instruction + user", len(inputs[0])) + } + if inputs[0][0].Role != schema.AgenticRoleTypeSystem || agenticMessageText(inputs[0][0]) != "system instruction" { + t.Fatalf("first agentic input = %#v", inputs[0][0]) + } + if inputs[0][1].Role != schema.AgenticRoleTypeUser || agenticMessageText(inputs[0][1]) != "classic input" { + t.Fatalf("second agentic input = %#v", inputs[0][1]) + } + + snapshot := trace.Snapshot() + if len(snapshot) != 2 || snapshot[0].Role != schema.System || snapshot[1].Role != schema.User { + t.Fatalf("trace snapshot = %#v, want classic system + user trace", snapshot) + } +} + +func TestNewEinoAgenticChatModelAgentAdapterPreservesTypedToolCallsForToolLayerRecovery(t *testing.T) { + t.Parallel() + ctx := context.Background() + fakeModel := &capturingAgenticChatModel{ + output: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call-1", + Name: "exec", + Arguments: `{"command":"` + strings.Repeat("x", 20000) + `"}`, + })}, + }, + } + agent, err := newEinoAgenticChatModelAgentAdapter(ctx, einoAgenticChatModelAgentConfig{ + Name: "agentic", + Description: "agentic adapter test", + Model: fakeModel, + Handlers: appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + phase: "agentic", + }), + }) + if err != nil { + t.Fatalf("newEinoAgenticChatModelAgentAdapter: %v", err) + } + iter := agent.Run(ctx, &adk.AgentInput{Messages: []*schema.Message{schema.UserMessage("run")}}) + var last *adk.AgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + t.Fatalf("agent event error: %v", ev.Err) + } + last = ev + } + if last == nil || last.Output == nil || last.Output.MessageOutput == nil { + t.Fatalf("last event = %#v, want message output", last) + } + msg := last.Output.MessageOutput.Message + if len(msg.ToolCalls) != 1 { + t.Fatalf("tool calls = %#v, want one tool call", msg.ToolCalls) + } + args := msg.ToolCalls[0].Function.Arguments + if !strings.Contains(args, strings.Repeat("x", 32)) || strings.Contains(args, modelOutputRecoveryKey) { + t.Fatalf("agentic tool args were unexpectedly rewritten: %q", args) + } +} + +func TestNewEinoAgenticChatModelAgentAdapterRequiresModel(t *testing.T) { + t.Parallel() + if _, err := newEinoAgenticChatModelAgentAdapter(context.Background(), einoAgenticChatModelAgentConfig{}); err == nil { + t.Fatal("expected missing model error") + } +} diff --git a/internal/multiagent/eino_agentic_chat_model_tail_middleware.go b/internal/multiagent/eino_agentic_chat_model_tail_middleware.go new file mode 100644 index 00000000..a82cd1ae --- /dev/null +++ b/internal/multiagent/eino_agentic_chat_model_tail_middleware.go @@ -0,0 +1,209 @@ +package multiagent + +import ( + "context" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// appendEinoAgenticChatModelTailMiddlewares appends protocol-neutral handlers for +// TypedChatModelAgent[*schema.AgenticMessage]. Classic ReAct history repair +// handlers stay on the schema.Message path because AgenticMessage has native +// content blocks for function calls/results. +func appendEinoAgenticChatModelTailMiddlewares( + handlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], + cfg einoChatModelTailConfig, +) []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + handlers = append(handlers, newAgenticSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase)) + handlers = append(handlers, newAgenticContinuationUserDedupMiddleware(cfg.logger, cfg.phase)) + if cfg.agenticSummarization != nil { + handlers = append(handlers, cfg.agenticSummarization) + } + if !cfg.skipTrace && cfg.trace != nil { + if capMw := newAgenticModelFacingTraceMiddleware(cfg.trace); capMw != nil { + handlers = append(handlers, capMw) + } + } + return handlers +} + +type agenticSystemMessageNormalizerMiddleware struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + logger *zap.Logger + phase string +} + +func newAgenticSystemMessageNormalizerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + return &agenticSystemMessageNormalizerMiddleware{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + logger: logger, + phase: phase, + } +} + +func (m *agenticSystemMessageNormalizerMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + mc *adk.TypedModelContext[*schema.AgenticMessage], +) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + _ = mc + if m == nil || state == nil || len(state.Messages) == 0 { + return ctx, state, nil + } + before := countAgenticSystemMessages(state.Messages) + if before <= 1 { + return ctx, state, nil + } + normalized := normalizeSingleLeadingAgenticSystemMessage(state.Messages) + if len(normalized) == len(state.Messages) && countAgenticSystemMessages(normalized) >= before { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Info("eino agentic system messages merged", + zap.String("phase", m.phase), + zap.Int("system_before", before), + zap.Int("system_after", countAgenticSystemMessages(normalized)), + zap.Int("messages_before", len(state.Messages)), + zap.Int("messages_after", len(normalized)), + ) + } + out := *state + out.Messages = normalized + return ctx, &out, nil +} + +type agenticContinuationUserDedupMiddleware struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + logger *zap.Logger + phase string +} + +func newAgenticContinuationUserDedupMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + return &agenticContinuationUserDedupMiddleware{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + logger: logger, + phase: phase, + } +} + +func (m *agenticContinuationUserDedupMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + mc *adk.TypedModelContext[*schema.AgenticMessage], +) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + _ = mc + if m == nil || state == nil || len(state.Messages) == 0 { + return ctx, state, nil + } + deduped, dropped := dedupAgenticContinuationUserMessages(state.Messages) + if dropped == 0 { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Info("eino agentic continuation user messages deduplicated", + zap.String("phase", m.phase), + zap.Int("dropped", dropped), + zap.Int("messages_before", len(state.Messages)), + zap.Int("messages_after", len(deduped)), + ) + } + out := *state + out.Messages = deduped + return ctx, &out, nil +} + +func countAgenticSystemMessages(msgs []*schema.AgenticMessage) int { + n := 0 + for _, msg := range msgs { + if msg != nil && msg.Role == schema.AgenticRoleTypeSystem { + n++ + } + } + return n +} + +func normalizeSingleLeadingAgenticSystemMessage(msgs []*schema.AgenticMessage) []*schema.AgenticMessage { + var systemParts []string + out := make([]*schema.AgenticMessage, 0, len(msgs)) + for _, msg := range msgs { + if msg == nil { + continue + } + if msg.Role == schema.AgenticRoleTypeSystem { + if text := strings.TrimSpace(agenticMessageText(msg)); text != "" { + systemParts = append(systemParts, text) + } + continue + } + out = append(out, msg) + } + if len(systemParts) == 0 { + return out + } + merged := schema.SystemAgenticMessage(strings.Join(systemParts, "\n\n")) + return append([]*schema.AgenticMessage{merged}, out...) +} + +func dedupAgenticContinuationUserMessages(msgs []*schema.AgenticMessage) ([]*schema.AgenticMessage, int) { + lastIdx := -1 + contCount := 0 + for i, msg := range msgs { + if !isAgenticContinuationUserMessage(msg) { + continue + } + contCount++ + lastIdx = i + } + if contCount <= 1 { + return msgs, 0 + } + out := make([]*schema.AgenticMessage, 0, len(msgs)-(contCount-1)) + dropped := 0 + for i, msg := range msgs { + if isAgenticContinuationUserMessage(msg) && i != lastIdx { + dropped++ + continue + } + out = append(out, msg) + } + return out, dropped +} + +func isAgenticContinuationUserMessage(msg *schema.AgenticMessage) bool { + if msg == nil || msg.Role != schema.AgenticRoleTypeUser { + return false + } + return strings.Contains(agenticMessageText(msg), continuationSessionMarker) +} + +func agenticMessageText(msg *schema.AgenticMessage) string { + if msg == nil { + return "" + } + var b strings.Builder + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + switch { + case block.UserInputText != nil: + if s := strings.TrimSpace(block.UserInputText.Text); s != "" { + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(s) + } + case block.AssistantGenText != nil: + if s := strings.TrimSpace(block.AssistantGenText.Text); s != "" { + if b.Len() > 0 { + b.WriteByte('\n') + } + b.WriteString(s) + } + } + } + return b.String() +} diff --git a/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go b/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go new file mode 100644 index 00000000..d020eb7d --- /dev/null +++ b/internal/multiagent/eino_agentic_chat_model_tail_middleware_test.go @@ -0,0 +1,112 @@ +package multiagent + +import ( + "context" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestAgenticSystemMessageNormalizerMiddlewareMergesDuplicates(t *testing.T) { + t.Parallel() + mw := newAgenticSystemMessageNormalizerMiddleware(nil, "test") + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.SystemAgenticMessage("first"), + schema.UserAgenticMessage("hello"), + schema.SystemAgenticMessage("second"), + }, + } + _, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + if out == state { + t.Fatal("expected rewritten state") + } + if got := countAgenticSystemMessages(out.Messages); got != 1 { + t.Fatalf("system messages = %d, want 1", got) + } + if out.Messages[0].Role != schema.AgenticRoleTypeSystem { + t.Fatalf("first role = %s, want system", out.Messages[0].Role) + } + text := agenticMessageText(out.Messages[0]) + if !strings.Contains(text, "first") || !strings.Contains(text, "second") { + t.Fatalf("merged system text = %q", text) + } + if len(out.Messages) != 2 || agenticMessageText(out.Messages[1]) != "hello" { + t.Fatalf("normalized messages = %#v", out.Messages) + } +} + +func TestAgenticContinuationUserDedupMiddlewareKeepsLatest(t *testing.T) { + t.Parallel() + mw := newAgenticContinuationUserDedupMiddleware(nil, "test") + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage(continuationSessionMarker + "\nold"), + schema.UserAgenticMessage("real user request"), + schema.UserAgenticMessage(continuationSessionMarker + "\nnew"), + }, + } + _, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + if out == state { + t.Fatal("expected rewritten state") + } + if len(out.Messages) != 2 { + t.Fatalf("messages = %d, want 2", len(out.Messages)) + } + if strings.Contains(agenticMessageText(out.Messages[0]), continuationSessionMarker) { + t.Fatalf("old continuation was not dropped: %#v", out.Messages) + } + if !strings.Contains(agenticMessageText(out.Messages[1]), "new") { + t.Fatalf("latest continuation not retained: %#v", out.Messages) + } +} + +func TestAgenticModelFacingTraceMiddlewareStoresClassicTrace(t *testing.T) { + t.Parallel() + holder := newModelFacingTraceHolder() + mw := newAgenticModelFacingTraceMiddleware(holder) + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.SystemAgenticMessage("instruction"), + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "answer"}), + }, + }, + }, + } + if _, _, err := mw.BeforeModelRewriteState(context.Background(), state, nil); err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + got := holder.Snapshot() + if len(got) != 2 { + t.Fatalf("trace len = %d, want 2", len(got)) + } + if got[0].Role != schema.System || got[0].Content != "instruction" { + t.Fatalf("system trace = %#v", got[0]) + } + if got[1].Role != schema.Assistant || got[1].Content != "answer" { + t.Fatalf("assistant trace = %#v", got[1]) + } +} + +func TestAppendEinoAgenticChatModelTailMiddlewares(t *testing.T) { + t.Parallel() + holder := newModelFacingTraceHolder() + handlers := appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + phase: "agentic", + trace: holder, + }) + if len(handlers) != 3 { + t.Fatalf("handlers = %d, want system + continuation + trace", len(handlers)) + } +} diff --git a/internal/multiagent/eino_agentic_event_adapter.go b/internal/multiagent/eino_agentic_event_adapter.go new file mode 100644 index 00000000..0b91c49d --- /dev/null +++ b/internal/multiagent/eino_agentic_event_adapter.go @@ -0,0 +1,106 @@ +package multiagent + +import ( + "io" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// adaptAgenticEventToEinoEvents converts typed AgenticMessage ADK events into +// the classic schema.Message events consumed by the existing SSE/MCP drain. +func adaptAgenticEventToEinoEvents(ev *adk.TypedAgentEvent[*schema.AgenticMessage]) []*adk.AgentEvent { + if ev == nil { + return nil + } + base := func(output *adk.AgentOutput) *adk.AgentEvent { + return &adk.AgentEvent{ + AgentName: ev.AgentName, + RunPath: append([]adk.RunStep(nil), ev.RunPath...), + Output: output, + Action: ev.Action, + Err: ev.Err, + } + } + if ev.Output == nil { + return []*adk.AgentEvent{base(nil)} + } + customized := ev.Output.CustomizedOutput + mv := ev.Output.MessageOutput + if mv == nil { + return []*adk.AgentEvent{base(&adk.AgentOutput{CustomizedOutput: customized})} + } + if mv.IsStreaming { + return []*adk.AgentEvent{base(&adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + IsStreaming: true, + MessageStream: agenticStreamToEinoStream(mv.MessageStream), + Role: agenticVariantRole(mv), + }, + CustomizedOutput: customized, + })} + } + + msgs := AgenticMessageToEino(mv.Message) + if len(msgs) == 0 { + return []*adk.AgentEvent{base(&adk.AgentOutput{CustomizedOutput: customized})} + } + out := make([]*adk.AgentEvent, 0, len(msgs)) + for i, msg := range msgs { + eventCustomized := any(nil) + if i == 0 { + eventCustomized = customized + } + out = append(out, base(&adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: msg, + Role: msg.Role, + ToolName: msg.ToolName, + }, + CustomizedOutput: eventCustomized, + })) + } + return out +} + +func agenticStreamToEinoStream(sr *schema.StreamReader[*schema.AgenticMessage]) *schema.StreamReader[*schema.Message] { + out, writer := schema.Pipe[*schema.Message](8) + go func() { + defer writer.Close() + if sr == nil { + return + } + defer sr.Close() + for { + chunk, err := sr.Recv() + if err != nil { + if err != io.EOF { + writer.Send(nil, err) + } + return + } + for _, msg := range AgenticMessageToEino(chunk) { + if msg != nil && writer.Send(msg, nil) { + return + } + } + } + }() + return out +} + +func agenticVariantRole(mv *adk.TypedMessageVariant[*schema.AgenticMessage]) schema.RoleType { + if mv == nil { + return schema.Assistant + } + switch mv.AgenticRole { + case schema.AgenticRoleTypeSystem: + return schema.System + case schema.AgenticRoleTypeUser: + // In Agentic ReAct output, user-role events from the graph are local + // FunctionToolResult messages emitted by AgenticToolsNode. + return schema.Tool + default: + return schema.Assistant + } +} diff --git a/internal/multiagent/eino_agentic_event_adapter_test.go b/internal/multiagent/eino_agentic_event_adapter_test.go new file mode 100644 index 00000000..e60bbbea --- /dev/null +++ b/internal/multiagent/eino_agentic_event_adapter_test.go @@ -0,0 +1,249 @@ +package multiagent + +import ( + "errors" + "io" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestAdaptAgenticEventToEinoEventsAssistantMessage(t *testing.T) { + usage := &schema.TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15} + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "agentic", + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ResponseMeta: &schema.AgenticResponseMeta{TokenUsage: usage}, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.Reasoning{Text: "think"}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "calling"}), + schema.NewContentBlock(&schema.FunctionToolCall{CallID: "call-1", Name: "scan", Arguments: `{"host":"127.0.0.1"}`}), + }, + }, + }, + CustomizedOutput: "custom", + }, + } + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 1 { + t.Fatalf("events = %d, want 1", len(got)) + } + mv := got[0].Output.MessageOutput + if got[0].AgentName != "agentic" || got[0].Output.CustomizedOutput != "custom" { + t.Fatalf("event metadata = %#v", got[0]) + } + if mv.Role != schema.Assistant || mv.Message.Role != schema.Assistant { + t.Fatalf("role = %q/%q, want assistant", mv.Role, mv.Message.Role) + } + if mv.Message.Content != "calling" || mv.Message.ReasoningContent != "think" { + t.Fatalf("message text = %#v", mv.Message) + } + if len(mv.Message.ToolCalls) != 1 || mv.Message.ToolCalls[0].ID != "call-1" || mv.Message.ToolCalls[0].Function.Name != "scan" { + t.Fatalf("tool calls = %#v", mv.Message.ToolCalls) + } + if mv.Message.ResponseMeta == nil || mv.Message.ResponseMeta.Usage != usage { + t.Fatalf("usage = %#v, want original usage", mv.Message.ResponseMeta) + } +} + +func TestAdaptAgenticEventToEinoEventsPureToolResult(t *testing.T) { + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{ + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-2", + Name: "execute", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: "done"}, + }}, + }), + }, + }, + }, + }, + } + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 1 { + t.Fatalf("events = %d, want 1", len(got)) + } + msg := got[0].Output.MessageOutput.Message + if got[0].Output.MessageOutput.Role != schema.Tool || msg.Role != schema.Tool || msg.ToolName != "execute" || msg.ToolCallID != "call-2" || msg.Content != "done" { + t.Fatalf("tool event = %#v message=%#v", got[0].Output.MessageOutput, msg) + } +} + +func TestAdaptAgenticEventToEinoEventsSplitsMixedToolResult(t *testing.T) { + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{ + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "text"}), + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-3", + Name: "grep", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: "match"}, + }}, + }), + }, + }, + }, + }, + } + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 2 { + t.Fatalf("events = %d, want assistant + tool", len(got)) + } + if got[0].Output.MessageOutput.Role != schema.Assistant || got[0].Output.MessageOutput.Message.Content != "text" { + t.Fatalf("assistant event = %#v", got[0].Output.MessageOutput) + } + if got[1].Output.MessageOutput.Role != schema.Tool || got[1].Output.MessageOutput.Message.ToolName != "grep" { + t.Fatalf("tool event = %#v", got[1].Output.MessageOutput) + } +} + +func TestAdaptAgenticEventToEinoEventsStreamingAssistant(t *testing.T) { + stream := schema.StreamReaderFromArray([]*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "hel"}), + }, + }, + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "lo"}), + }, + }, + }) + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{ + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: stream, + AgenticRole: schema.AgenticRoleTypeAssistant, + }, + }, + } + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 1 { + t.Fatalf("events = %d, want 1", len(got)) + } + mv := got[0].Output.MessageOutput + if !mv.IsStreaming || mv.Role != schema.Assistant { + t.Fatalf("stream variant = %#v", mv) + } + first, err := mv.MessageStream.Recv() + if err != nil || first.Content != "hel" { + t.Fatalf("first = %#v err=%v", first, err) + } + second, err := mv.MessageStream.Recv() + if err != nil || second.Content != "lo" { + t.Fatalf("second = %#v err=%v", second, err) + } + _, err = mv.MessageStream.Recv() + if !errors.Is(err, io.EOF) { + t.Fatalf("final err = %v, want EOF", err) + } +} + +func TestAdaptAgenticStreamingToolResultFeedsClassicToolResultHandler(t *testing.T) { + stream := schema.StreamReaderFromArray([]*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-agentic-stream", + Name: "execute", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: "partial "}, + }}, + }), + }, + }, + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-agentic-stream", + Name: "execute", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: "done"}, + }}, + }), + }, + }, + }) + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "agentic", + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: stream, + AgenticRole: schema.AgenticRoleTypeUser, + }, + }, + } + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 1 || got[0].Output == nil || got[0].Output.MessageOutput == nil { + t.Fatalf("events = %#v", got) + } + mv := got[0].Output.MessageOutput + if !mv.IsStreaming || mv.Role != schema.Tool { + t.Fatalf("streaming variant = %#v, want tool stream", mv) + } + + var event map[string]interface{} + runMessages := newEinoRunMessageAccumulator(nil) + emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{ + ConversationID: "conv-agentic", + Progress: func(eventType, _ string, data interface{}) { + if eventType == "tool_result" { + event, _ = data.(map[string]interface{}) + } + }, + }) + handler := newEinoToolResultEventHandler(einoToolResultEventHandlerConfig{ + RunMessages: runMessages, + Emitter: emitter, + }) + if !handler.HandleStreaming(mv, "agentic") { + t.Fatal("agentic streaming tool result was not handled") + } + if event["toolName"] != "execute" || event["toolCallId"] != "call-agentic-stream" || event["result"] != "partial done" { + t.Fatalf("tool result event = %#v", event) + } + msgs := runMessages.Messages() + if len(msgs) != 1 || msgs[0].ToolName != "execute" || msgs[0].ToolCallID != "call-agentic-stream" || msgs[0].Content != "partial done" { + t.Fatalf("run messages = %#v", msgs) + } +} + +func TestAdaptAgenticEventToEinoEventsPreservesErrorOnlyEvent(t *testing.T) { + wantErr := errors.New("boom") + ev := &adk.TypedAgentEvent[*schema.AgenticMessage]{AgentName: "agentic", Err: wantErr} + + got := adaptAgenticEventToEinoEvents(ev) + if len(got) != 1 || got[0].AgentName != "agentic" || !errors.Is(got[0].Err, wantErr) { + t.Fatalf("events = %#v", got) + } +} diff --git a/internal/multiagent/eino_agentic_message.go b/internal/multiagent/eino_agentic_message.go new file mode 100644 index 00000000..d74003e7 --- /dev/null +++ b/internal/multiagent/eino_agentic_message.go @@ -0,0 +1,190 @@ +package multiagent + +import ( + "strings" + + "github.com/cloudwego/eino/schema" +) + +// EinoMessagesToAgentic converts the project's current ADK message history to +// Eino's native AgenticMessage shape. It intentionally covers the text, +// reasoning, function tool-call, and function tool-result channels used by the +// agent runtime today; unsupported multimodal/provider-specific fields stay in +// schema.Message until a real AgenticModel backend is wired. +func EinoMessagesToAgentic(msgs []*schema.Message) []*schema.AgenticMessage { + if len(msgs) == 0 { + return nil + } + out := make([]*schema.AgenticMessage, 0, len(msgs)) + for _, msg := range msgs { + if msg == nil { + continue + } + out = append(out, EinoMessageToAgentic(msg)) + } + return out +} + +func EinoMessageToAgentic(msg *schema.Message) *schema.AgenticMessage { + if msg == nil { + return nil + } + out := &schema.AgenticMessage{ + Role: messageRoleToAgentic(msg.Role), + Extra: cloneAnyMap(msg.Extra), + } + if msg.ResponseMeta != nil { + out.ResponseMeta = &schema.AgenticResponseMeta{TokenUsage: msg.ResponseMeta.Usage} + } + if text := strings.TrimSpace(msg.ReasoningContent); text != "" { + out.ContentBlocks = append(out.ContentBlocks, schema.NewContentBlock(&schema.Reasoning{Text: msg.ReasoningContent})) + } + switch msg.Role { + case schema.Assistant: + if msg.Content != "" { + out.ContentBlocks = append(out.ContentBlocks, schema.NewContentBlock(&schema.AssistantGenText{Text: msg.Content})) + } + for _, tc := range msg.ToolCalls { + out.ContentBlocks = append(out.ContentBlocks, schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: tc.ID, + Name: tc.Function.Name, + Arguments: tc.Function.Arguments, + })) + } + case schema.Tool: + out.Role = schema.AgenticRoleTypeUser + out.ContentBlocks = append(out.ContentBlocks, schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: msg.ToolCallID, + Name: msg.ToolName, + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: msg.Content}, + }}, + })) + default: + if msg.Content != "" { + out.ContentBlocks = append(out.ContentBlocks, schema.NewContentBlock(&schema.UserInputText{Text: msg.Content})) + } + } + return out +} + +// AgenticMessagesToEino converts AgenticMessage values back into the classic +// schema.Message form used by the existing ADK event drain and persistence code. +func AgenticMessagesToEino(msgs []*schema.AgenticMessage) []*schema.Message { + if len(msgs) == 0 { + return nil + } + out := make([]*schema.Message, 0, len(msgs)) + for _, msg := range msgs { + if msg == nil { + continue + } + out = append(out, AgenticMessageToEino(msg)...) + } + return out +} + +func AgenticMessageToEino(msg *schema.AgenticMessage) []*schema.Message { + if msg == nil { + return nil + } + base := &schema.Message{ + Role: agenticRoleToMessage(msg.Role), + Extra: cloneAnyMap(msg.Extra), + } + if msg.ResponseMeta != nil { + base.ResponseMeta = &schema.ResponseMeta{Usage: msg.ResponseMeta.TokenUsage} + } + var toolResults []*schema.Message + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + switch { + case block.Reasoning != nil: + base.ReasoningContent += block.Reasoning.Text + case block.UserInputText != nil: + base.Content += block.UserInputText.Text + case block.AssistantGenText != nil: + base.Role = schema.Assistant + base.Content += block.AssistantGenText.Text + case block.FunctionToolCall != nil: + var index *int + if block.StreamingMeta != nil { + i := block.StreamingMeta.Index + index = &i + } + base.Role = schema.Assistant + base.ToolCalls = append(base.ToolCalls, schema.ToolCall{ + Index: index, + ID: block.FunctionToolCall.CallID, + Type: "function", + Function: schema.FunctionCall{ + Name: block.FunctionToolCall.Name, + Arguments: block.FunctionToolCall.Arguments, + }, + }) + case block.FunctionToolResult != nil: + toolResults = append(toolResults, functionToolResultToMessage(block.FunctionToolResult)) + } + } + if len(toolResults) > 0 && base.Content == "" && base.ReasoningContent == "" && len(base.ToolCalls) == 0 { + return toolResults + } + out := []*schema.Message{base} + out = append(out, toolResults...) + return out +} + +func messageRoleToAgentic(role schema.RoleType) schema.AgenticRoleType { + switch role { + case schema.System: + return schema.AgenticRoleTypeSystem + case schema.Assistant: + return schema.AgenticRoleTypeAssistant + default: + return schema.AgenticRoleTypeUser + } +} + +func agenticRoleToMessage(role schema.AgenticRoleType) schema.RoleType { + switch role { + case schema.AgenticRoleTypeSystem: + return schema.System + case schema.AgenticRoleTypeAssistant: + return schema.Assistant + default: + return schema.User + } +} + +func functionToolResultToMessage(result *schema.FunctionToolResult) *schema.Message { + if result == nil { + return nil + } + parts := make([]string, 0, len(result.Content)) + for _, block := range result.Content { + if block == nil || block.Text == nil { + continue + } + parts = append(parts, block.Text.Text) + } + return &schema.Message{ + Role: schema.Tool, + Content: strings.Join(parts, ""), + ToolCallID: result.CallID, + ToolName: result.Name, + } +} + +func cloneAnyMap(in map[string]any) map[string]any { + if len(in) == 0 { + return nil + } + out := make(map[string]any, len(in)) + for k, v := range in { + out[k] = v + } + return out +} diff --git a/internal/multiagent/eino_agentic_message_test.go b/internal/multiagent/eino_agentic_message_test.go new file mode 100644 index 00000000..9fa89ab6 --- /dev/null +++ b/internal/multiagent/eino_agentic_message_test.go @@ -0,0 +1,154 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/schema" +) + +func TestEinoMessageToAgenticPreservesAssistantToolCalls(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + Content: "I will scan it.", + ReasoningContent: "Need enumerate first.", + ToolCalls: []schema.ToolCall{{ + ID: "call-1", + Type: "function", + Function: schema.FunctionCall{ + Name: "nmap", + Arguments: `{"target":"127.0.0.1"}`, + }, + }}, + Extra: map[string]any{"trace": "kept"}, + } + + got := EinoMessageToAgentic(msg) + if got.Role != schema.AgenticRoleTypeAssistant { + t.Fatalf("role = %q, want assistant", got.Role) + } + if len(got.ContentBlocks) != 3 { + t.Fatalf("blocks = %d, want 3", len(got.ContentBlocks)) + } + if got.ContentBlocks[0].Reasoning == nil || got.ContentBlocks[0].Reasoning.Text != msg.ReasoningContent { + t.Fatalf("reasoning block = %#v", got.ContentBlocks[0]) + } + if got.ContentBlocks[1].AssistantGenText == nil || got.ContentBlocks[1].AssistantGenText.Text != msg.Content { + t.Fatalf("assistant text block = %#v", got.ContentBlocks[1]) + } + call := got.ContentBlocks[2].FunctionToolCall + if call == nil || call.CallID != "call-1" || call.Name != "nmap" || call.Arguments != `{"target":"127.0.0.1"}` { + t.Fatalf("tool call block = %#v", got.ContentBlocks[2]) + } + if got.Extra["trace"] != "kept" { + t.Fatalf("extra = %#v", got.Extra) + } +} + +func TestEinoMessageToAgenticMapsToolResultAsUserFunctionResult(t *testing.T) { + msg := &schema.Message{ + Role: schema.Tool, + Content: "22/tcp open ssh", + ToolCallID: "call-ssh", + ToolName: "nmap", + } + + got := EinoMessageToAgentic(msg) + if got.Role != schema.AgenticRoleTypeUser { + t.Fatalf("role = %q, want user", got.Role) + } + if len(got.ContentBlocks) != 1 || got.ContentBlocks[0].FunctionToolResult == nil { + t.Fatalf("blocks = %#v", got.ContentBlocks) + } + result := got.ContentBlocks[0].FunctionToolResult + if result.CallID != "call-ssh" || result.Name != "nmap" { + t.Fatalf("tool result metadata = %#v", result) + } + if len(result.Content) != 1 || result.Content[0].Text == nil || result.Content[0].Text.Text != "22/tcp open ssh" { + t.Fatalf("tool result content = %#v", result.Content) + } +} + +func TestAgenticMessageToEinoPreservesAssistantBlocks(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.Reasoning{Text: "Think first."}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "Calling scanner."}), + schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call-2", + Name: "scan", + Arguments: `{"host":"example.com"}`, + }), + }, + } + + got := AgenticMessageToEino(msg) + if len(got) != 1 { + t.Fatalf("messages = %d, want 1", len(got)) + } + if got[0].Role != schema.Assistant || got[0].Content != "Calling scanner." || got[0].ReasoningContent != "Think first." { + t.Fatalf("assistant message = %#v", got[0]) + } + if len(got[0].ToolCalls) != 1 || got[0].ToolCalls[0].ID != "call-2" || got[0].ToolCalls[0].Function.Name != "scan" { + t.Fatalf("tool calls = %#v", got[0].ToolCalls) + } +} + +func TestAgenticMessageToEinoSplitsPureToolResult(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-3", + Name: "execute", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: "done"}, + }}, + }), + }, + } + + got := AgenticMessageToEino(msg) + if len(got) != 1 { + t.Fatalf("messages = %d, want 1", len(got)) + } + if got[0].Role != schema.Tool || got[0].ToolCallID != "call-3" || got[0].ToolName != "execute" || got[0].Content != "done" { + t.Fatalf("tool message = %#v", got[0]) + } +} + +func TestEinoAgenticRoundTripForSupportedFields(t *testing.T) { + msgs := []*schema.Message{ + schema.SystemMessage("system"), + schema.UserMessage("user"), + { + Role: schema.Assistant, + Content: "assistant", + ToolCalls: []schema.ToolCall{{ + ID: "call-4", + Type: "function", + Function: schema.FunctionCall{Name: "grep", Arguments: `{"q":"token"}`}, + }}, + }, + { + Role: schema.Tool, + Content: "match", + ToolCallID: "call-4", + ToolName: "grep", + }, + } + + got := AgenticMessagesToEino(EinoMessagesToAgentic(msgs)) + if len(got) != len(msgs) { + t.Fatalf("round trip messages = %d, want %d: %#v", len(got), len(msgs), got) + } + for i := range msgs { + if got[i].Role != msgs[i].Role || got[i].Content != msgs[i].Content || got[i].ToolCallID != msgs[i].ToolCallID || got[i].ToolName != msgs[i].ToolName { + t.Fatalf("message[%d] = %#v, want %#v", i, got[i], msgs[i]) + } + if len(got[i].ToolCalls) != len(msgs[i].ToolCalls) { + t.Fatalf("message[%d] tool calls = %#v, want %#v", i, got[i].ToolCalls, msgs[i].ToolCalls) + } + } +} diff --git a/internal/multiagent/eino_agentic_model_gate.go b/internal/multiagent/eino_agentic_model_gate.go new file mode 100644 index 00000000..1b4e8db1 --- /dev/null +++ b/internal/multiagent/eino_agentic_model_gate.go @@ -0,0 +1,109 @@ +package multiagent + +import ( + "context" + "strings" + + "github.com/cloudwego/eino/components/model" + "go.uber.org/zap" +) + +type einoAgenticModelFactory func(context.Context) (model.AgenticModel, error) + +type einoAgenticRuntimeSupport struct { + TypedRunner bool + Streaming bool + CancelMonitoring bool + ModelRetry bool + ModelFailover bool + ToolResultObservation bool + MCPExecutionAudit bool +} + +type einoAgenticModelGate struct { + Ready bool + Reason string + Missing []string +} + +// Eino v0.9.14 wires AgenticMessage through the same generic TypedRunner, +// stream cancel monitoring, model retry, and model failover wrappers used by +// schema.Message. Keep this matrix explicit so future upgrades are audited +// deliberately instead of flipping the AgenticModel path by accident. +func einoAgenticRuntimeSupportV0914() einoAgenticRuntimeSupport { + return einoAgenticRuntimeSupport{ + TypedRunner: true, + Streaming: true, + CancelMonitoring: true, + ModelRetry: true, + ModelFailover: true, + ToolResultObservation: true, + MCPExecutionAudit: true, + } +} + +func evaluateEinoAgenticModelGate(factory einoAgenticModelFactory, support einoAgenticRuntimeSupport) einoAgenticModelGate { + missing := make([]string, 0, 8) + if factory == nil { + missing = append(missing, "model.AgenticModel backend") + } else { + if m, err := factory(context.Background()); err != nil || m == nil { + missing = append(missing, "model.AgenticModel backend") + } + } + if !support.TypedRunner { + missing = append(missing, "adk.TypedRunner[*schema.AgenticMessage]") + } + if !support.Streaming { + missing = append(missing, "AgenticMessage streaming") + } + if !support.CancelMonitoring { + missing = append(missing, "AgenticMessage model-stream cancel monitoring") + } + if !support.ModelRetry { + missing = append(missing, "AgenticMessage ModelRetry") + } + if !support.ModelFailover { + missing = append(missing, "AgenticMessage ModelFailover") + } + if !support.ToolResultObservation { + missing = append(missing, "AgenticMessage tool-result observation") + } + if !support.MCPExecutionAudit { + missing = append(missing, "AgenticMessage MCP execution audit") + } + if len(missing) == 0 { + return einoAgenticModelGate{Ready: true, Reason: "ready"} + } + return einoAgenticModelGate{ + Reason: "agentic_model_not_ready: " + strings.Join(missing, ", "), + Missing: missing, + } +} + +func logEinoAgenticModelGate(logger *zap.Logger, scope, orchestration string, gate einoAgenticModelGate) { + if logger == nil { + return + } + fields := []zap.Field{ + zap.String("scope", scope), + zap.String("orchestration", orchestration), + zap.Bool("ready", gate.Ready), + zap.String("reason", gate.Reason), + zap.Strings("missing", gate.Missing), + } + if gate.Ready { + logger.Info("eino agentic model gate ready", fields...) + return + } + logger.Info("eino agentic model gate disabled", fields...) +} + +func agenticTextModelFactory(m model.AgenticModel) einoAgenticModelFactory { + if m == nil { + return nil + } + return func(context.Context) (model.AgenticModel, error) { + return m, nil + } +} diff --git a/internal/multiagent/eino_agentic_model_gate_test.go b/internal/multiagent/eino_agentic_model_gate_test.go new file mode 100644 index 00000000..78ec18c1 --- /dev/null +++ b/internal/multiagent/eino_agentic_model_gate_test.go @@ -0,0 +1,93 @@ +package multiagent + +import ( + "context" + "errors" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type fakeAgenticGateModel struct{} + +func (m *fakeAgenticGateModel) Generate(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.AgenticMessage, error) { + return &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant}, nil +} + +func (m *fakeAgenticGateModel) Stream(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + return schema.StreamReaderFromArray([]*schema.AgenticMessage{{Role: schema.AgenticRoleTypeAssistant}}), nil +} + +func TestEinoAgenticModelGateV0914WaitsOnlyForBackend(t *testing.T) { + gate := evaluateEinoAgenticModelGate(nil, einoAgenticRuntimeSupportV0914()) + + if gate.Ready { + t.Fatal("v0.9.14 gate should stay disabled without an AgenticModel backend") + } + if !containsString(gate.Missing, "model.AgenticModel backend") { + t.Fatalf("missing = %#v, want backend reason", gate.Missing) + } + for _, unexpected := range []string{ + "AgenticMessage model-stream cancel monitoring", + "AgenticMessage ModelRetry", + "AgenticMessage ModelFailover", + "AgenticMessage tool-result observation", + "AgenticMessage MCP execution audit", + } { + if containsString(gate.Missing, unexpected) { + t.Fatalf("missing = %#v, should not include %q for v0.9.14 runtime support", gate.Missing, unexpected) + } + } +} + +func TestEinoAgenticModelGateV0914ReadyWithBackend(t *testing.T) { + gate := evaluateEinoAgenticModelGate(agenticTextModelFactory(&fakeAgenticGateModel{}), einoAgenticRuntimeSupportV0914()) + + if !gate.Ready { + t.Fatalf("gate = %#v, want ready when v0.9.14 runtime support has a backend", gate) + } + if gate.Reason != "ready" || len(gate.Missing) != 0 { + t.Fatalf("gate details = %#v", gate) + } +} + +func TestEinoAgenticModelGateReadyWhenBackendAndRuntimeParityExist(t *testing.T) { + gate := evaluateEinoAgenticModelGate(agenticTextModelFactory(&fakeAgenticGateModel{}), einoAgenticRuntimeSupport{ + TypedRunner: true, + Streaming: true, + CancelMonitoring: true, + ModelRetry: true, + ModelFailover: true, + ToolResultObservation: true, + MCPExecutionAudit: true, + }) + + if !gate.Ready { + t.Fatalf("gate = %#v, want ready", gate) + } + if gate.Reason != "ready" || len(gate.Missing) != 0 { + t.Fatalf("gate details = %#v", gate) + } +} + +func TestEinoAgenticModelGateTreatsFactoryErrorAsMissingBackend(t *testing.T) { + gate := evaluateEinoAgenticModelGate(func(context.Context) (model.AgenticModel, error) { + return nil, errors.New("not implemented") + }, einoAgenticRuntimeSupport{ + TypedRunner: true, + Streaming: true, + CancelMonitoring: true, + ModelRetry: true, + ModelFailover: true, + ToolResultObservation: true, + MCPExecutionAudit: true, + }) + + if gate.Ready { + t.Fatal("factory error should disable gate") + } + if !containsString(gate.Missing, "model.AgenticModel backend") { + t.Fatalf("missing = %#v, want backend reason", gate.Missing) + } +} diff --git a/internal/multiagent/eino_agentic_summarize.go b/internal/multiagent/eino_agentic_summarize.go new file mode 100644 index 00000000..d4f1e45b --- /dev/null +++ b/internal/multiagent/eino_agentic_summarize.go @@ -0,0 +1,278 @@ +package multiagent + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/middlewares/summarization" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// newEinoAgenticSummarizationMiddleware wires the project's domain-specific +// compaction policy into Eino's native typed AgenticMessage summarization. +func newEinoAgenticSummarizationMiddleware( + ctx context.Context, + summaryModel model.BaseModel[*schema.AgenticMessage], + appCfg *config.Config, + mwCfg *config.MultiAgentEinoMiddlewareConfig, + conversationID string, + db *database.DB, + projectID string, + logger *zap.Logger, +) (adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], error) { + if summaryModel == nil || appCfg == nil { + return nil, fmt.Errorf("multiagent: agentic summarization 需要 model 与配置") + } + maxTotal := appCfg.OpenAI.MaxTotalTokens + if maxTotal <= 0 { + maxTotal = 120000 + } + triggerRatio := 0.8 + emitInternalEvents := true + outputReserve := config.DefaultSummarizationOutputReserveTokens + userLedgerMaxRunes := config.DefaultSummarizationUserIntentLedgerMaxRunes + userLedgerEntryMaxRunes := config.DefaultSummarizationUserIntentLedgerEntryMaxRunes + toolMaxBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective() + if mwCfg != nil { + triggerRatio = mwCfg.SummarizationTriggerRatioEffective() + emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective() + outputReserve = mwCfg.SummarizationOutputReserveTokensEffective() + userLedgerMaxRunes = mwCfg.SummarizationUserIntentLedgerMaxRunesEffective() + userLedgerEntryMaxRunes = mwCfg.SummarizationUserIntentLedgerEntryMaxRunesEffective() + toolMaxBytes = mwCfg.ReductionMaxLengthForTruncEffective() + } + + ledgerWindowCap := modelFacingRuneBudget(maxTotal, 0.20) + userLedgerMaxRunes = minPositiveInt(userLedgerMaxRunes, ledgerWindowCap) + userLedgerEntryMaxRunes = minPositiveInt(userLedgerEntryMaxRunes, userLedgerMaxRunes) + trigger := int(float64(maxTotal) * triggerRatio) + if trigger < 4096 { + trigger = maxTotal + if trigger < 4096 { + trigger = 4096 + } + } + modelName := strings.TrimSpace(appCfg.OpenAI.Model) + if modelName == "" { + modelName = "gpt-4o" + } + classicTokenCounter := einoSummarizationTokenCounter(modelName) + agenticTokenCounter := func(ctx context.Context, input *summarization.TypedTokenCounterInput[*schema.AgenticMessage]) (int, error) { + if input == nil { + return 0, nil + } + return classicTokenCounter(ctx, &summarization.TokenCounterInput{ + Messages: AgenticMessagesToEino(input.Messages), + Tools: input.Tools, + }) + } + recentTrailMax := trigger / 4 + if recentTrailMax < 2048 { + recentTrailMax = 2048 + } + if recentTrailMax > trigger/2 { + recentTrailMax = trigger / 2 + } + summaryInputMax := trigger - outputReserve + if summaryInputMax < 4096 { + summaryInputMax = trigger * 80 / 100 + } + if summaryInputMax < 4096 { + summaryInputMax = 4096 + } + + transcriptPath := "" + if conv := strings.TrimSpace(conversationID); conv != "" { + baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization") + if dbPath := strings.TrimSpace(appCfg.Database.Path); dbPath != "" { + baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization") + } + base := baseRoot + if abs, err := filepath.Abs(base); err == nil { + base = abs + } + if mkErr := os.MkdirAll(base, 0o755); mkErr == nil { + transcriptPath = filepath.Join(base, "transcript.txt") + } + } + + retryPolicy := einoTransientRunRetryPolicyFromMW(mwCfg) + retryMax := retryPolicy.maxAttempts + var summaryOverflowRetries int + summaryModelOpts := []model.Option{ + einoopenai.WithMaxCompletionTokens(outputReserve), + } + + mw, err := summarization.NewTyped[*schema.AgenticMessage](ctx, &summarization.TypedConfig[*schema.AgenticMessage]{ + Model: summaryModel, + ModelOptions: summaryModelOpts, + GenModelInput: func(ctx context.Context, sysInstruction, userInstruction *schema.AgenticMessage, originalMsgs []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) { + classicOriginal := AgenticMessagesToEino(originalMsgs) + if transcriptPath != "" && len(classicOriginal) > 0 { + if werr := writeSummarizationTranscript(transcriptPath, classicOriginal); werr != nil && logger != nil { + logger.Warn("eino agentic summarization transcript preflight 写入失败", + zap.String("path", transcriptPath), zap.Error(werr)) + } + } + budget := summaryInputMax + aggressive := summaryOverflowRetries > 0 + if aggressive { + budget = summaryInputMax * 70 / 100 + if budget < 4096 { + budget = 4096 + } + } + input, dropped, berr := buildBudgetedSummarizationModelInput( + ctx, + agenticInstructionToClassic(sysInstruction, schema.System), + agenticInstructionToClassic(userInstruction, schema.User), + classicOriginal, + classicTokenCounter, + budget, + summarizationInputBudgetOpts{ + toolMaxBytes: toolMaxBytes, + spillRef: transcriptPath, + aggressive: aggressive, + }, + ) + if logger != nil && (berr != nil || dropped > 0 || aggressive) { + fields := []zap.Field{ + zap.Int("max_input_tokens", budget), + zap.Int("trigger_context_tokens", trigger), + zap.Int("output_reserve_tokens", outputReserve), + zap.Int("dropped_rounds", dropped), + zap.Bool("aggressive", aggressive), + } + if berr != nil { + fields = append(fields, zap.Error(berr)) + logger.Warn("eino agentic summarization input budget failed", fields...) + } else { + logger.Info("eino agentic summarization input bounded", fields...) + } + } + return EinoMessagesToAgentic(input), berr + }, + Trigger: &summarization.TriggerCondition{ + ContextTokens: trigger, + }, + TokenCounter: agenticTokenCounter, + UserInstruction: einoSummarizeUserInstruction, + EmitInternalEvents: emitInternalEvents, + TranscriptFilePath: transcriptPath, + Retry: &summarization.TypedRetryConfig[*schema.AgenticMessage]{ + MaxRetries: &retryMax, + ShouldRetry: func(_ context.Context, _ *schema.AgenticMessage, err error) bool { + if isEinoContextOverflowError(err) && summaryOverflowRetries < 1 { + summaryOverflowRetries++ + if logger != nil { + logger.Warn("eino agentic summarization context overflow, retrying with aggressive compaction", + zap.Error(err), + ) + } + return true + } + retry := isEinoTransientRunError(err) + if retry && logger != nil { + logger.Warn("eino agentic summarization generate transient error, will retry if attempts remain", + zap.Error(err), + zap.Int("max_retries", retryMax), + ) + } + return retry + }, + }, + Finalize: func(ctx context.Context, originalMessages []*schema.AgenticMessage, summary *schema.AgenticMessage) ([]*schema.AgenticMessage, error) { + classicOriginal := AgenticMessagesToEino(originalMessages) + classicSummary := agenticSummaryToClassicMessage(summary) + if classicSummary == nil { + return nil, fmt.Errorf("agentic summarization returned empty summary") + } + compactionMessages := stripOriginalUserIntentLedgerFromMessages(classicOriginal) + defaultFinalized, derr := summarization.DefaultFinalize(ctx, compactionMessages, classicSummary) + if derr != nil { + return nil, derr + } + if len(defaultFinalized) == 0 { + return nil, fmt.Errorf("agentic summarization default finalize returned no messages") + } + summaryMsg := appendTranscriptPathToSummarizationMessage(defaultFinalized[len(defaultFinalized)-1], transcriptPath) + summaryMsg = stripAnalysisFromSummarizationMessage(summaryMsg) + userLedger := buildOriginalUserIntentLedgerMessage(classicOriginal, userLedgerMaxRunes, userLedgerEntryMaxRunes) + out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, compactionMessages, summaryMsg, classicTokenCounter, recentTrailMax) + if ferr != nil { + return nil, ferr + } + out = mergeMessageIntoLeadingSystem(out, userLedger) + if appCfg != nil { + out = refreshFactIndexInMessages(out, db, projectID, appCfg.Project, logger) + } + return EinoMessagesToAgentic(out), nil + }, + Callback: func(ctx context.Context, before, after adk.TypedChatModelAgentState[*schema.AgenticMessage]) error { + classicBefore := AgenticMessagesToEino(before.Messages) + classicAfter := AgenticMessagesToEino(after.Messages) + if transcriptPath != "" && len(classicBefore) > 0 { + if werr := writeSummarizationTranscript(transcriptPath, classicBefore); werr != nil && logger != nil { + logger.Warn("eino agentic summarization transcript 写入失败", + zap.String("path", transcriptPath), + zap.Error(werr), + ) + } + } + if logger != nil { + beforeTokens, _ := classicTokenCounter(ctx, &summarization.TokenCounterInput{Messages: classicBefore}) + afterTokens, _ := classicTokenCounter(ctx, &summarization.TokenCounterInput{Messages: classicAfter}) + logger.Info("eino agentic summarization 已压缩上下文", + zap.Int("messages_before", len(before.Messages)), + zap.Int("messages_after", len(after.Messages)), + zap.Int("tokens_before_estimated", beforeTokens), + zap.Int("tokens_after_estimated", afterTokens), + zap.Int("max_total_tokens", maxTotal), + zap.Int("trigger_context_tokens", trigger), + zap.String("transcript_file", transcriptPath), + ) + } + return nil + }, + }) + if err != nil { + return nil, fmt.Errorf("summarization.NewTyped[AgenticMessage]: %w", err) + } + return mw, nil +} + +func agenticInstructionToClassic(msg *schema.AgenticMessage, fallbackRole schema.RoleType) *schema.Message { + msgs := AgenticMessageToEino(msg) + if len(msgs) > 0 && msgs[0] != nil { + return msgs[0] + } + return &schema.Message{Role: fallbackRole} +} + +func agenticSummaryToClassicMessage(msg *schema.AgenticMessage) *schema.Message { + msgs := AgenticMessageToEino(msg) + for _, m := range msgs { + if m == nil { + continue + } + if m.Role == schema.Assistant || strings.TrimSpace(m.Content) != "" || m.ReasoningContent != "" { + if m.Role != schema.Assistant { + cp := *m + cp.Role = schema.Assistant + return &cp + } + return m + } + } + return nil +} diff --git a/internal/multiagent/eino_agentic_summarize_test.go b/internal/multiagent/eino_agentic_summarize_test.go new file mode 100644 index 00000000..4472582c --- /dev/null +++ b/internal/multiagent/eino_agentic_summarize_test.go @@ -0,0 +1,210 @@ +package multiagent + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestNewEinoAgenticSummarizationMiddlewareCompactsWithNativeTypedMiddleware(t *testing.T) { + t.Parallel() + ctx := context.Background() + emit := false + summaryModel := &capturingAgenticChatModel{ + output: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: `检查历史 + +## 1. 授权范围与约束 +- 仅测试 example.com + +## 7. 当前进度、策略决策与下一步 +- 继续验证 SQL 注入路径 +`})}, + }, + } + appCfg := &config.Config{} + appCfg.OpenAI.Model = "gpt-4o" + appCfg.OpenAI.MaxTotalTokens = 5000 + appCfg.Database.Path = filepath.Join(t.TempDir(), "cyberstrike.db") + mwCfg := &config.MultiAgentEinoMiddlewareConfig{ + SummarizationEmitInternalEvents: &emit, + SummarizationOutputReserveTokens: 1024, + } + + mw, err := newEinoAgenticSummarizationMiddleware(ctx, summaryModel, appCfg, mwCfg, "conv-agentic", nil, "", nil) + if err != nil { + t.Fatalf("newEinoAgenticSummarizationMiddleware: %v", err) + } + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.SystemAgenticMessage("system root"), + schema.UserAgenticMessage("授权范围 example.com\n" + strings.Repeat("历史扫描输出 ", 12000)), + agenticAssistantTextMessage("已记录范围"), + schema.UserAgenticMessage("继续验证 SQL 注入路径"), + }, + } + + _, after, err := mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + inputs := summaryModel.snapshotInputs() + if len(inputs) != 1 || len(inputs[0]) == 0 { + t.Fatalf("summary model inputs = %#v, want one typed AgenticMessage call", inputs) + } + if after == nil { + t.Fatal("after state is nil") + } + classicAfter := AgenticMessagesToEino(after.Messages) + joined := joinClassicMessageContent(classicAfter) + if strings.Contains(joined, "") { + t.Fatalf("analysis block leaked into compacted context: %s", joined) + } + for _, want := range []string{"继续验证 SQL 注入路径", "原始用户输入与约束账本", "完整的对话记录位于"} { + if !strings.Contains(joined, want) { + t.Fatalf("compacted context missing %q:\n%s", want, joined) + } + } +} + +func TestEinoAgenticChatModelAgentCompactsContextBeforeBusinessModel(t *testing.T) { + t.Parallel() + ctx := context.Background() + emit := false + summaryModel := &capturingAgenticChatModel{ + output: agenticAssistantTextMessage(`internal scratchpad + +## 1. 授权范围与约束 +- 仅测试 example.com + +## 7. 当前进度、策略决策与下一步 +- 继续验证 SQL 注入路径 +`), + } + businessModel := &capturingAgenticChatModel{ + output: agenticAssistantTextMessage("business answer after compaction"), + } + appCfg := &config.Config{} + appCfg.OpenAI.Model = "gpt-4o" + appCfg.OpenAI.MaxTotalTokens = 5000 + appCfg.Database.Path = filepath.Join(t.TempDir(), "cyberstrike.db") + mwCfg := &config.MultiAgentEinoMiddlewareConfig{ + SummarizationEmitInternalEvents: &emit, + SummarizationOutputReserveTokens: 1024, + } + sumMw, err := newEinoAgenticSummarizationMiddleware(ctx, summaryModel, appCfg, mwCfg, "conv-agentic-e2e", nil, "", nil) + if err != nil { + t.Fatalf("newEinoAgenticSummarizationMiddleware: %v", err) + } + trace := newModelFacingTraceHolder() + agent, err := newEinoAgenticChatModelAgentAdapter(ctx, einoAgenticChatModelAgentConfig{ + Name: "agentic", + Description: "agentic compaction e2e test", + Instruction: "system root", + Model: businessModel, + Handlers: appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + phase: "agentic", + agenticSummarization: sumMw, + trace: trace, + }), + }) + if err != nil { + t.Fatalf("newEinoAgenticChatModelAgentAdapter: %v", err) + } + + rawHistory := "授权范围 example.com\n" + strings.Repeat("原始扫描输出SHOULD_NOT_REACH_BUSINESS_MODEL ", 12000) + iter := agent.Run(ctx, &adk.AgentInput{ + Messages: []*schema.Message{ + schema.UserMessage(rawHistory), + schema.AssistantMessage("已记录范围", nil), + schema.UserMessage("继续验证 SQL 注入路径"), + }, + }) + var last *adk.AgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + t.Fatalf("agent event error: %v", ev.Err) + } + last = ev + } + if last == nil || last.Output == nil || last.Output.MessageOutput == nil { + t.Fatalf("last event = %#v, want message output", last) + } + if got := last.Output.MessageOutput.Message.Content; got != "business answer after compaction" { + t.Fatalf("business output = %q", got) + } + + if inputs := summaryModel.snapshotInputs(); len(inputs) != 1 { + t.Fatalf("summary model calls = %d, want 1", len(inputs)) + } + businessInputs := businessModel.snapshotInputs() + if len(businessInputs) != 1 { + t.Fatalf("business model calls = %d, want 1", len(businessInputs)) + } + finalClassicInput := AgenticMessagesToEino(businessInputs[0]) + joined := joinClassicMessageContent(finalClassicInput) + for _, want := range []string{"继续验证 SQL 注入路径", "原始用户输入与约束账本", "完整的对话记录位于"} { + if !strings.Contains(joined, want) { + t.Fatalf("business model input missing %q:\n%s", want, joined) + } + } + if strings.Contains(joined, "") { + t.Fatalf("analysis leaked to business model input:\n%s", joined) + } + if strings.Count(joined, "原始扫描输出SHOULD_NOT_REACH_BUSINESS_MODEL") > 3 { + t.Fatalf("raw oversized history leaked to business model input, count=%d", strings.Count(joined, "原始扫描输出SHOULD_NOT_REACH_BUSINESS_MODEL")) + } + traceJoined := joinClassicMessageContent(trace.Snapshot()) + if !strings.Contains(traceJoined, "继续验证 SQL 注入路径") || strings.Count(traceJoined, "原始扫描输出SHOULD_NOT_REACH_BUSINESS_MODEL") > 3 { + t.Fatalf("model-facing trace not compacted:\n%s", traceJoined) + } +} + +func TestAppendEinoAgenticChatModelTailMiddlewaresIncludesTypedSummarization(t *testing.T) { + t.Parallel() + mw := newAgenticSystemMessageNormalizerMiddleware(nil, "summary") + handlers := appendEinoAgenticChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + agenticSummarization: mw, + skipTrace: true, + }) + found := false + for _, h := range handlers { + if h == mw { + found = true + break + } + } + if !found { + t.Fatal("agentic summarization middleware was not appended") + } +} + +func agenticAssistantTextMessage(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: text})}, + } +} + +func joinClassicMessageContent(msgs []*schema.Message) string { + var b strings.Builder + for _, msg := range msgs { + if msg == nil { + continue + } + b.WriteString(msg.Content) + b.WriteByte('\n') + } + return b.String() +} diff --git a/internal/multiagent/eino_agentic_tool_calling_adapter.go b/internal/multiagent/eino_agentic_tool_calling_adapter.go new file mode 100644 index 00000000..50d913f0 --- /dev/null +++ b/internal/multiagent/eino_agentic_tool_calling_adapter.go @@ -0,0 +1,118 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// agenticToolCallingChatModelAdapter lets Eino's classic plan_execute +// Planner/Replanner consume a native AgenticModel without translating the HTTP +// protocol. Only the in-memory Eino message and option shapes are adapted. +type agenticToolCallingChatModelAdapter struct { + model model.AgenticModel + tools []*schema.ToolInfo +} + +func newAgenticToolCallingChatModelAdapter(agenticModel model.AgenticModel) model.ToolCallingChatModel { + return &agenticToolCallingChatModelAdapter{model: agenticModel} +} + +func (m *agenticToolCallingChatModelAdapter) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + if m == nil || m.model == nil { + return nil, fmt.Errorf("agentic tool-calling adapter: model is nil") + } + clonedTools := append([]*schema.ToolInfo(nil), tools...) + return &agenticToolCallingChatModelAdapter{model: m.model, tools: clonedTools}, nil +} + +func (m *agenticToolCallingChatModelAdapter) Generate( + ctx context.Context, + input []*schema.Message, + opts ...model.Option, +) (*schema.Message, error) { + if m == nil || m.model == nil { + return nil, fmt.Errorf("agentic tool-calling adapter: model is nil") + } + out, err := m.model.Generate(ctx, EinoMessagesToAgentic(input), m.agenticOptions(opts...)...) + if err != nil { + return nil, err + } + converted := AgenticMessageToEino(out) + if len(converted) == 0 { + return nil, fmt.Errorf("agentic tool-calling adapter: model returned no message") + } + return converted[0], nil +} + +func (m *agenticToolCallingChatModelAdapter) Stream( + ctx context.Context, + input []*schema.Message, + opts ...model.Option, +) (*schema.StreamReader[*schema.Message], error) { + if m == nil || m.model == nil { + return nil, fmt.Errorf("agentic tool-calling adapter: model is nil") + } + stream, err := m.model.Stream(ctx, EinoMessagesToAgentic(input), m.agenticOptions(opts...)...) + if err != nil { + return nil, err + } + return agenticStreamToEinoStream(stream), nil +} + +func (m *agenticToolCallingChatModelAdapter) agenticOptions(opts ...model.Option) []model.Option { + common := model.GetCommonOptions(&model.Options{ + Tools: append([]*schema.ToolInfo(nil), m.tools...), + }, opts...) + out := make([]model.Option, 0, 8) + if common.Temperature != nil { + out = append(out, model.WithTemperature(*common.Temperature)) + } + if common.Model != nil { + out = append(out, model.WithModel(*common.Model)) + } + if common.TopP != nil { + out = append(out, model.WithTopP(*common.TopP)) + } + if common.MaxTokens != nil { + out = append(out, model.WithMaxTokens(*common.MaxTokens)) + } + if len(common.Stop) > 0 { + out = append(out, model.WithStop(common.Stop)) + } + if common.Tools != nil { + out = append(out, model.WithTools(common.Tools)) + } + if common.AgenticToolChoice != nil { + out = append(out, model.WithAgenticToolChoice(common.AgenticToolChoice)) + } else if common.ToolChoice != nil { + out = append(out, model.WithAgenticToolChoice(classicToolChoiceToAgentic( + *common.ToolChoice, + common.AllowedToolNames, + ))) + } + return out +} + +func classicToolChoiceToAgentic(choice schema.ToolChoice, allowedNames []string) *schema.AgenticToolChoice { + allowed := make([]*schema.AllowedTool, 0, len(allowedNames)) + for _, name := range allowedNames { + if name != "" { + allowed = append(allowed, &schema.AllowedTool{FunctionName: name}) + } + } + out := &schema.AgenticToolChoice{Type: choice} + switch choice { + case schema.ToolChoiceAllowed: + if len(allowed) > 0 { + out.Allowed = &schema.AgenticAllowedToolChoice{Tools: allowed} + } + case schema.ToolChoiceForced: + if len(allowed) > 0 { + out.Forced = &schema.AgenticForcedToolChoice{Tools: allowed} + } + } + return out +} diff --git a/internal/multiagent/eino_agentic_tool_calling_adapter_test.go b/internal/multiagent/eino_agentic_tool_calling_adapter_test.go new file mode 100644 index 00000000..71a0c150 --- /dev/null +++ b/internal/multiagent/eino_agentic_tool_calling_adapter_test.go @@ -0,0 +1,118 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type capturingAgenticToolCallingModel struct { + input []*schema.AgenticMessage + options *model.Options +} + +func (m *capturingAgenticToolCallingModel) Generate( + _ context.Context, + input []*schema.AgenticMessage, + opts ...model.Option, +) (*schema.AgenticMessage, error) { + m.input = input + m.options = model.GetCommonOptions(nil, opts...) + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call-1", + Name: "emit_plan", + Arguments: `{"steps":["inspect"]}`, + }), + }, + }, nil +} + +func (m *capturingAgenticToolCallingModel) Stream( + context.Context, + []*schema.AgenticMessage, + ...model.Option, +) (*schema.StreamReader[*schema.AgenticMessage], error) { + return schema.StreamReaderFromArray([]*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlockChunk(&schema.FunctionToolCall{ + CallID: "call-1", + Name: "emit_plan", + Arguments: `{"steps":[`, + }, &schema.StreamingMeta{Index: 0}), + }, + }, + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlockChunk(&schema.FunctionToolCall{ + Arguments: `"inspect"]}`, + }, &schema.StreamingMeta{Index: 0}), + }, + }, + }), nil +} + +func TestAgenticToolCallingAdapterConvertsForcedToolChoice(t *testing.T) { + t.Parallel() + native := &capturingAgenticToolCallingModel{} + adapter, err := newAgenticToolCallingChatModelAdapter(native).WithTools([]*schema.ToolInfo{{ + Name: "emit_plan", + Desc: "emit a structured plan", + }}) + if err != nil { + t.Fatalf("WithTools: %v", err) + } + + out, err := adapter.Generate( + context.Background(), + []*schema.Message{schema.UserMessage("plan this task")}, + model.WithToolChoice(schema.ToolChoiceForced), + ) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if len(native.input) != 1 || native.input[0].Role != schema.AgenticRoleTypeUser { + t.Fatalf("native input = %#v", native.input) + } + if native.options == nil || len(native.options.Tools) != 1 || native.options.Tools[0].Name != "emit_plan" { + t.Fatalf("native tools = %#v", native.options) + } + if native.options.AgenticToolChoice == nil || native.options.AgenticToolChoice.Type != schema.ToolChoiceForced { + t.Fatalf("agentic tool choice = %#v", native.options.AgenticToolChoice) + } + if len(out.ToolCalls) != 1 || out.ToolCalls[0].Function.Name != "emit_plan" { + t.Fatalf("classic output = %#v", out) + } +} + +func TestAgenticToolCallingAdapterPreservesStreamingToolCallIndex(t *testing.T) { + t.Parallel() + adapter := newAgenticToolCallingChatModelAdapter(&capturingAgenticToolCallingModel{}) + stream, err := adapter.Stream(context.Background(), []*schema.Message{ + schema.UserMessage("plan this task"), + }) + if err != nil { + t.Fatalf("Stream: %v", err) + } + out, err := schema.ConcatMessageStream(stream) + if err != nil { + t.Fatalf("ConcatMessageStream: %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("tool calls = %#v, want one merged call", out.ToolCalls) + } + call := out.ToolCalls[0] + if call.Index == nil || *call.Index != 0 { + t.Fatalf("tool call index = %#v", call.Index) + } + if call.Function.Name != "emit_plan" || call.Function.Arguments != `{"steps":["inspect"]}` { + t.Fatalf("merged tool call = %#v", call) + } +} diff --git a/internal/multiagent/eino_assistant_output_accumulator.go b/internal/multiagent/eino_assistant_output_accumulator.go new file mode 100644 index 00000000..1ca53128 --- /dev/null +++ b/internal/multiagent/eino_assistant_output_accumulator.go @@ -0,0 +1,42 @@ +package multiagent + +import "strings" + +type einoAssistantOutputAccumulator struct { + orchMode string + lastAssistant string + lastPlanExecuteExecutor string +} + +func newEinoAssistantOutputAccumulator(orchMode string) *einoAssistantOutputAccumulator { + return &einoAssistantOutputAccumulator{orchMode: orchMode} +} + +func (a *einoAssistantOutputAccumulator) RecordMainAssistant(agentName, content string) bool { + if a == nil { + return false + } + content = strings.TrimSpace(content) + if content == "" { + return false + } + a.lastAssistant = content + if a.orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(agentName), "executor") { + a.lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(content) + } + return true +} + +func (a *einoAssistantOutputAccumulator) LastAssistant() string { + if a == nil { + return "" + } + return a.lastAssistant +} + +func (a *einoAssistantOutputAccumulator) LastPlanExecuteExecutor() string { + if a == nil { + return "" + } + return a.lastPlanExecuteExecutor +} diff --git a/internal/multiagent/eino_assistant_output_accumulator_test.go b/internal/multiagent/eino_assistant_output_accumulator_test.go new file mode 100644 index 00000000..27102b06 --- /dev/null +++ b/internal/multiagent/eino_assistant_output_accumulator_test.go @@ -0,0 +1,52 @@ +package multiagent + +import "testing" + +func TestEinoAssistantOutputAccumulatorRecordsMainAssistant(t *testing.T) { + acc := newEinoAssistantOutputAccumulator("deep") + if acc.RecordMainAssistant("lead", " hello ") != true { + t.Fatal("expected record") + } + if got := acc.LastAssistant(); got != "hello" { + t.Fatalf("last assistant = %q, want hello", got) + } + if got := acc.LastPlanExecuteExecutor(); got != "" { + t.Fatalf("plan execute executor = %q, want empty", got) + } + if acc.RecordMainAssistant("lead", " ") { + t.Fatal("blank content should not record") + } + if got := acc.LastAssistant(); got != "hello" { + t.Fatalf("blank content changed last assistant to %q", got) + } +} + +func TestEinoAssistantOutputAccumulatorPlanExecuteExecutor(t *testing.T) { + acc := newEinoAssistantOutputAccumulator("plan_execute") + raw := `{"response":"给用户看的正文","scratchpad":"internal"}` + acc.RecordMainAssistant("executor", raw) + + if got := acc.LastAssistant(); got != raw { + t.Fatalf("last assistant = %q, want raw", got) + } + if got := acc.LastPlanExecuteExecutor(); got != "给用户看的正文" { + t.Fatalf("executor output = %q", got) + } + acc.RecordMainAssistant("planner", "planner note") + if got := acc.LastAssistant(); got != "planner note" { + t.Fatalf("last assistant after planner = %q", got) + } + if got := acc.LastPlanExecuteExecutor(); got != "给用户看的正文" { + t.Fatalf("planner should not overwrite executor output, got %q", got) + } +} + +func TestEinoAssistantOutputAccumulatorNilSafe(t *testing.T) { + var acc *einoAssistantOutputAccumulator + if acc.RecordMainAssistant("agent", "hello") { + t.Fatal("nil accumulator should not record") + } + if acc.LastAssistant() != "" || acc.LastPlanExecuteExecutor() != "" { + t.Fatal("nil accumulator should return empty values") + } +} diff --git a/internal/multiagent/eino_assistant_stream_event_handler.go b/internal/multiagent/eino_assistant_stream_event_handler.go new file mode 100644 index 00000000..978d946f --- /dev/null +++ b/internal/multiagent/eino_assistant_stream_event_handler.go @@ -0,0 +1,166 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +type einoAssistantStreamEventHandlerConfig struct { + Context context.Context + ConversationID string + OrchMode string + Progress func(eventType, message string, data interface{}) + Logger *zap.Logger + SnapshotMCPIDs func() []string + StreamsMainAssistant func(agent string) bool + EinoRoleTag func(agent string) string + RunProgress *einoRunProgressTracker + StdoutSuppressor *einoExecuteStdoutSuppressor + AssistantOutput *einoAssistantOutputAccumulator + RunMessages *einoRunMessageAccumulator + Usage *einoRunUsageAccumulator + ToolCallCompletion *einoStreamToolCallCompletionHandler + NextMainStreamID func() string + NextReasoningStreamID func() string + NextSubAgentReplyStreamID func() string +} + +type einoAssistantStreamEventHandler struct { + ctx context.Context + conversationID string + orchMode string + progress func(eventType, message string, data interface{}) + logger *zap.Logger + snapshotMCPIDs func() []string + streamsMainAssistant func(agent string) bool + einoRoleTag func(agent string) string + runProgress *einoRunProgressTracker + stdoutSuppressor *einoExecuteStdoutSuppressor + assistantOutput *einoAssistantOutputAccumulator + runMessages *einoRunMessageAccumulator + usage *einoRunUsageAccumulator + toolCallCompletion *einoStreamToolCallCompletionHandler + nextMainStreamID func() string + nextReasoningStreamID func() string + nextSubAgentReplyStreamID func() string +} + +func newEinoAssistantStreamEventHandler(cfg einoAssistantStreamEventHandlerConfig) *einoAssistantStreamEventHandler { + if cfg.Context == nil { + cfg.Context = context.Background() + } + if cfg.SnapshotMCPIDs == nil { + cfg.SnapshotMCPIDs = func() []string { return nil } + } + if cfg.StreamsMainAssistant == nil { + cfg.StreamsMainAssistant = func(string) bool { return true } + } + if cfg.EinoRoleTag == nil { + cfg.EinoRoleTag = func(string) string { return "" } + } + if cfg.NextMainStreamID == nil { + cfg.NextMainStreamID = func() string { return "eino-main" } + } + if cfg.NextReasoningStreamID == nil { + cfg.NextReasoningStreamID = func() string { return "eino-reasoning" } + } + if cfg.NextSubAgentReplyStreamID == nil { + cfg.NextSubAgentReplyStreamID = func() string { return "eino-sub-reply" } + } + return &einoAssistantStreamEventHandler{ + ctx: cfg.Context, + conversationID: cfg.ConversationID, + orchMode: cfg.OrchMode, + progress: cfg.Progress, + logger: cfg.Logger, + snapshotMCPIDs: cfg.SnapshotMCPIDs, + streamsMainAssistant: cfg.StreamsMainAssistant, + einoRoleTag: cfg.EinoRoleTag, + runProgress: cfg.RunProgress, + stdoutSuppressor: cfg.StdoutSuppressor, + assistantOutput: cfg.AssistantOutput, + runMessages: cfg.RunMessages, + usage: cfg.Usage, + toolCallCompletion: cfg.ToolCallCompletion, + nextMainStreamID: cfg.NextMainStreamID, + nextReasoningStreamID: cfg.NextReasoningStreamID, + nextSubAgentReplyStreamID: cfg.NextSubAgentReplyStreamID, + } +} + +func (h *einoAssistantStreamEventHandler) Handle(mv *adk.MessageVariant, agentName string) (handled bool, recvErr error) { + if h == nil || mv == nil || !mv.IsStreaming || mv.MessageStream == nil || mv.Role == schema.Tool { + return false, nil + } + mainStreamID := h.nextMainStreamID() + mainEmitter := newEinoMainResponseStreamEmitter( + h.conversationID, h.orchMode, agentName, mainStreamID, h.mainIteration(agentName), h.progress, h.snapshotMCPIDs, + ) + reasoningEmitter := newEinoReasoningStreamEmitter( + h.conversationID, + h.orchMode, + agentName, + h.einoRoleTag(agentName), + h.progress, + h.nextReasoningStreamID, + ) + var toolStreamFragments []schema.ToolCall + var streamUsage *schema.TokenUsage + subReplyEmitter := newEinoSubAgentReplyEmitter( + h.conversationID, + agentName, + h.progress, + h.nextSubAgentReplyStreamID, + ) + mainAssistantStream := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{ + AgentName: agentName, + Emitter: mainEmitter, + StdoutSuppressor: h.stdoutSuppressor, + AssistantOutput: h.assistantOutput, + RunMessages: h.runMessages, + }) + recvErr = recvEinoSchemaMessageStreamWithContext(h.ctx, mv.MessageStream, 8, func(chunk *schema.Message) { + reasoningEmitter.EmitDelta(chunk.ReasoningContent) + if chunk.Content != "" { + if h.streamsMainAssistant(agentName) { + mainAssistantStream.EmitDelta(chunk.Content) + } else if !h.streamsMainAssistant(agentName) { + subReplyEmitter.EmitDelta(chunk.Content) + } + } + if len(chunk.ToolCalls) > 0 { + toolStreamFragments = append(toolStreamFragments, chunk.ToolCalls...) + } + if chunk.ResponseMeta != nil && chunk.ResponseMeta.Usage != nil { + streamUsage = maxEinoTokenUsage(streamUsage, chunk.ResponseMeta.Usage) + } + }) + if recvErr != nil && !isEinoVoluntaryCancelErr(recvErr) && h.logger != nil { + h.logger.Warn("eino stream recv error, flushing incomplete stream", + zap.Error(recvErr), + zap.String("agent", agentName), + zap.Int("toolFragments", len(toolStreamFragments))) + } + reasoningEmitter.Finish() + if h.streamsMainAssistant(agentName) { + mainAssistantStream.Finish() + } + subReplyEmitter.Finish() + if h.toolCallCompletion != nil { + h.toolCallCompletion.Complete(toolStreamFragments, agentName) + } + if h.usage != nil { + h.usage.AddUsage(streamUsage) + } + return true, recvErr +} + +func (h *einoAssistantStreamEventHandler) mainIteration(agentName string) int { + if h == nil || h.runProgress == nil { + return 0 + } + return h.runProgress.MainIteration(agentName) +} diff --git a/internal/multiagent/eino_assistant_stream_event_handler_test.go b/internal/multiagent/eino_assistant_stream_event_handler_test.go new file mode 100644 index 00000000..05fecf1c --- /dev/null +++ b/internal/multiagent/eino_assistant_stream_event_handler_test.go @@ -0,0 +1,148 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestEinoAssistantStreamEventHandlerHandlesMainAssistantStream(t *testing.T) { + var events []string + runMessages := newEinoRunMessageAccumulator(nil) + assistantOutput := newEinoAssistantOutputAccumulator("deep") + usage := newEinoRunUsageAccumulator() + handler := newEinoAssistantStreamEventHandler(einoAssistantStreamEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + RunMessages: runMessages, + Usage: usage, + AssistantOutput: assistantOutput, + StreamsMainAssistant: func(agent string) bool { return agent == "lead" }, + EinoRoleTag: func(string) string { return "orchestrator" }, + NextMainStreamID: func() string { return "main-stream-1" }, + Progress: func(eventType, _ string, _ interface{}) { + events = append(events, eventType) + }, + }) + mv := &adk.MessageVariant{ + IsStreaming: true, + Role: schema.Assistant, + MessageStream: schema.StreamReaderFromArray([]*schema.Message{ + {Role: schema.Assistant, Content: "he", ResponseMeta: &schema.ResponseMeta{Usage: &schema.TokenUsage{PromptTokens: 10, CompletionTokens: 1, TotalTokens: 11}}}, + {Role: schema.Assistant, Content: "hello", ResponseMeta: &schema.ResponseMeta{Usage: &schema.TokenUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}}}, + }), + } + + handled, err := handler.Handle(mv, "lead") + if !handled || err != nil { + t.Fatalf("handled=%v err=%v", handled, err) + } + if assistantOutput.LastAssistant() != "hello" { + t.Fatalf("last assistant = %q", assistantOutput.LastAssistant()) + } + if msgs := runMessages.Messages(); len(msgs) != 1 || msgs[0].Content != "hello" { + t.Fatalf("run messages = %#v", msgs) + } + if got := usage.Summary(); got.ModelCalls != 1 || got.PromptTokens != 10 || got.CompletionTokens != 5 || got.TotalTokens != 15 { + t.Fatalf("usage = %#v, want one stream model call", got) + } + if !containsString(events, "response_start") || !containsString(events, "response_delta") { + t.Fatalf("events = %#v, want response stream events", events) + } +} + +func TestEinoAssistantStreamEventHandlerHandlesSubAgentStream(t *testing.T) { + var events []string + runMessages := newEinoRunMessageAccumulator(nil) + assistantOutput := newEinoAssistantOutputAccumulator("deep") + handler := newEinoAssistantStreamEventHandler(einoAssistantStreamEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + RunMessages: runMessages, + AssistantOutput: assistantOutput, + StreamsMainAssistant: func(agent string) bool { return agent == "lead" }, + EinoRoleTag: func(string) string { return "sub" }, + NextSubAgentReplyStreamID: func() string { + return "sub-stream-1" + }, + Progress: func(eventType, _ string, _ interface{}) { + events = append(events, eventType) + }, + }) + mv := &adk.MessageVariant{ + IsStreaming: true, + Role: schema.Assistant, + MessageStream: schema.StreamReaderFromArray([]*schema.Message{{Role: schema.Assistant, Content: "sub reply"}}), + } + + handled, err := handler.Handle(mv, "worker") + if !handled || err != nil { + t.Fatalf("handled=%v err=%v", handled, err) + } + if len(runMessages.Messages()) != 0 { + t.Fatalf("sub stream should not append main run text, got %#v", runMessages.Messages()) + } + if assistantOutput.LastAssistant() != "" { + t.Fatalf("sub stream should not record main assistant, got %q", assistantOutput.LastAssistant()) + } + if !containsString(events, "eino_agent_reply_stream_start") || + !containsString(events, "eino_agent_reply_stream_delta") || + !containsString(events, "eino_agent_reply_stream_end") { + t.Fatalf("events = %#v, want sub reply stream events", events) + } +} + +func TestEinoAssistantStreamEventHandlerCompletesToolFragments(t *testing.T) { + idx := 0 + var events []string + runMessages := newEinoRunMessageAccumulator(nil) + runProgress := newEinoRunProgressTracker( + "deep", "lead", "conv-1", + func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + func(agent string) bool { return agent == "lead" }, + nil, + ) + completion := newEinoStreamToolCallCompletionHandler(einoStreamToolCallCompletionHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + Progress: func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + RunProgress: runProgress, + RunMessages: runMessages, + }) + handler := newEinoAssistantStreamEventHandler(einoAssistantStreamEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + RunMessages: runMessages, + StreamsMainAssistant: func(string) bool { return true }, + ToolCallCompletion: completion, + }) + mv := &adk.MessageVariant{ + IsStreaming: true, + Role: schema.Assistant, + MessageStream: schema.StreamReaderFromArray([]*schema.Message{ + {Role: schema.Assistant, ToolCalls: []schema.ToolCall{{ID: "call-1", Index: &idx, Type: "function", Function: schema.FunctionCall{Name: "execute", Arguments: `{"command":`}}}}, + {Role: schema.Assistant, ToolCalls: []schema.ToolCall{{Index: &idx, Function: schema.FunctionCall{Arguments: `"pwd"}`}}}}, + }), + } + + handled, err := handler.Handle(mv, "lead") + if !handled || err != nil { + t.Fatalf("handled=%v err=%v", handled, err) + } + msgs := runMessages.Messages() + if len(msgs) != 1 || len(msgs[0].ToolCalls) != 1 || msgs[0].ToolCalls[0].Function.Arguments != `{"command":"pwd"}` { + t.Fatalf("run messages = %#v, want merged tool call", msgs) + } + if !containsString(events, "tool_call") { + t.Fatalf("events = %#v, want tool_call", events) + } +} + +func TestEinoAssistantStreamEventHandlerIgnoresToolStream(t *testing.T) { + handler := newEinoAssistantStreamEventHandler(einoAssistantStreamEventHandlerConfig{}) + handled, err := handler.Handle(&adk.MessageVariant{IsStreaming: true, Role: schema.Tool, MessageStream: schema.StreamReaderFromArray([]*schema.Message{})}, "lead") + if handled || err != nil { + t.Fatalf("handled=%v err=%v, want ignored", handled, err) + } +} diff --git a/internal/multiagent/eino_chat_model_tail_middleware.go b/internal/multiagent/eino_chat_model_tail_middleware.go new file mode 100644 index 00000000..66352887 --- /dev/null +++ b/internal/multiagent/eino_chat_model_tail_middleware.go @@ -0,0 +1,78 @@ +package multiagent + +import ( + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// einoChatModelTailConfig configures middleware appended after reduction/skill/plantask +// and immediately before each ChatModel invocation pipeline completes. +// +// Order (best practice): +// 1. system merge — accurate token count for summarization +// 2. continuation user dedup — drop stale session-resume injections +// 3. malformed tool-call arguments repair +// 4. pre-summarization tool-call/result reconciliation +// 5. summarization +// 6. soft model-input budget (warn/compact only, never fail locally) +// 7. final malformed tool-call arguments repair +// 8. final tool-call/result reconciliation +// 9. orphan tool prune (defense in depth) +// 10. malformed tool_search history repair +// 11. telemetry +// 12. model-facing trace snapshot +type einoChatModelTailConfig struct { + logger *zap.Logger + phase string + summarization adk.ChatModelAgentMiddleware + agenticSummarization adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + modelName string + maxTotalTokens int + toolMaxBytes int + conversationID string + trace *modelFacingTraceHolder + middlewareConfig *config.MultiAgentEinoMiddlewareConfig + skipOrphanPruner bool + skipTelemetry bool + skipTrace bool +} + +func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware, cfg einoChatModelTailConfig) []adk.ChatModelAgentMiddleware { + handlers = append(handlers, newSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase)) + handlers = append(handlers, newContinuationUserDedupMiddleware(cfg.logger, cfg.phase)) + handlers = append(handlers, newToolCallArgumentsSanitizerMiddleware(cfg.logger, cfg.phase+"_pre_summarization")) + if cfg.summarization != nil { + // Summarization invokes the model internally, so its input needs the same + // structural guarantee as the agent's final model call. + handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization")) + handlers = append(handlers, cfg.summarization) + } + handlers = append(handlers, newModelInputSoftBudgetMiddleware(cfg.maxTotalTokens, cfg.toolMaxBytes, cfg.modelName, cfg.logger, cfg.phase)) + handlers = append(handlers, newToolCallArgumentsSanitizerMiddleware(cfg.logger, cfg.phase)) + handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase)) + if !cfg.skipOrphanPruner { + handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase)) + } + handlers = append(handlers, newToolSearchResultSanitizerMiddleware(cfg.logger, cfg.phase)) + if !cfg.skipTelemetry { + if teleMw := newEinoModelInputTelemetryMiddleware(cfg.logger, cfg.modelName, cfg.conversationID, cfg.phase); teleMw != nil { + handlers = append(handlers, teleMw) + } + } + if !cfg.skipTrace && cfg.trace != nil { + if capMw := newModelFacingTraceMiddleware(cfg.trace); capMw != nil { + handlers = append(handlers, capMw) + } + } + return handlers +} + +func toolMaxBytesFromMW(mwCfg *config.MultiAgentEinoMiddlewareConfig) int { + if mwCfg != nil { + return mwCfg.ReductionMaxLengthForTruncEffective() + } + return config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective() +} diff --git a/internal/multiagent/eino_checkpoint.go b/internal/multiagent/eino_checkpoint.go new file mode 100644 index 00000000..569c698c --- /dev/null +++ b/internal/multiagent/eino_checkpoint.go @@ -0,0 +1,68 @@ +package multiagent + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" +) + +// fileCheckPointStore implements adk.CheckPointStore with one file per checkpoint id. +type fileCheckPointStore struct { + dir string +} + +func newFileCheckPointStore(baseDir string) (*fileCheckPointStore, error) { + if strings.TrimSpace(baseDir) == "" { + return nil, fmt.Errorf("checkpoint base dir empty") + } + abs, err := filepath.Abs(baseDir) + if err != nil { + return nil, err + } + if err := os.MkdirAll(abs, 0o755); err != nil { + return nil, err + } + return &fileCheckPointStore{dir: abs}, nil +} + +func (s *fileCheckPointStore) path(id string) (string, error) { + id = strings.TrimSpace(id) + if id == "" { + return "", fmt.Errorf("checkpoint id empty") + } + if strings.ContainsAny(id, `/\`) { + return "", fmt.Errorf("invalid checkpoint id") + } + return filepath.Join(s.dir, id+".ckpt"), nil +} + +func (s *fileCheckPointStore) Get(ctx context.Context, checkPointID string) ([]byte, bool, error) { + _ = ctx + p, err := s.path(checkPointID) + if err != nil { + return nil, false, err + } + b, err := os.ReadFile(p) + if err != nil { + if os.IsNotExist(err) { + return nil, false, nil + } + return nil, false, err + } + return b, true, nil +} + +func (s *fileCheckPointStore) Set(ctx context.Context, checkPointID string, checkPoint []byte) error { + _ = ctx + p, err := s.path(checkPointID) + if err != nil { + return err + } + tmp := p + ".tmp" + if err := os.WriteFile(tmp, checkPoint, 0o600); err != nil { + return err + } + return os.Rename(tmp, p) +} diff --git a/internal/multiagent/eino_checkpoint_resume_handler.go b/internal/multiagent/eino_checkpoint_resume_handler.go new file mode 100644 index 00000000..cb2e1f8b --- /dev/null +++ b/internal/multiagent/eino_checkpoint_resume_handler.go @@ -0,0 +1,71 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "go.uber.org/zap" +) + +type einoCheckpointResumeHandlerConfig struct { + Context context.Context + ConversationID string + OrchMode string + Progress func(eventType, message string, data interface{}) + Logger *zap.Logger + Store *fileCheckPointStore + CheckPointID string + Resume func(checkPointID string) (*adk.AsyncIterator[*adk.AgentEvent], error) +} + +type einoCheckpointResumeHandler struct { + cfg einoCheckpointResumeHandlerConfig +} + +func newEinoCheckpointResumeHandler(cfg einoCheckpointResumeHandlerConfig) *einoCheckpointResumeHandler { + if cfg.Context == nil { + cfg.Context = context.Background() + } + return &einoCheckpointResumeHandler{cfg: cfg} +} + +func (h *einoCheckpointResumeHandler) TryResume() *adk.AsyncIterator[*adk.AgentEvent] { + if h == nil || h.cfg.Store == nil || h.cfg.CheckPointID == "" || h.cfg.Resume == nil { + return nil + } + if _, existed, err := h.cfg.Store.Get(h.cfg.Context, h.cfg.CheckPointID); err != nil { + if h.cfg.Logger != nil { + h.cfg.Logger.Warn("eino checkpoint preflight get failed", zap.String("checkPointID", h.cfg.CheckPointID), zap.Error(err)) + } + return nil + } else if !existed { + return nil + } + h.emitProgress("检测到断点,正在从中断节点恢复执行...") + if h.cfg.Logger != nil { + h.cfg.Logger.Info("eino runner: resume from checkpoint", zap.String("checkPointID", h.cfg.CheckPointID)) + } + iter, err := h.cfg.Resume(h.cfg.CheckPointID) + if err == nil { + return iter + } + if h.cfg.Logger != nil { + h.cfg.Logger.Warn("eino runner: resume failed, fallback to fresh run", + zap.String("checkPointID", h.cfg.CheckPointID), + zap.Error(err)) + } + h.emitProgress("断点恢复失败,已回退为全新执行。") + return nil +} + +func (h *einoCheckpointResumeHandler) emitProgress(message string) { + if h == nil || h.cfg.Progress == nil { + return + } + h.cfg.Progress("progress", message, map[string]interface{}{ + "conversationId": h.cfg.ConversationID, + "source": "eino", + "orchestration": h.cfg.OrchMode, + "checkPointID": h.cfg.CheckPointID, + }) +} diff --git a/internal/multiagent/eino_checkpoint_resume_handler_test.go b/internal/multiagent/eino_checkpoint_resume_handler_test.go new file mode 100644 index 00000000..33b1372f --- /dev/null +++ b/internal/multiagent/eino_checkpoint_resume_handler_test.go @@ -0,0 +1,138 @@ +package multiagent + +import ( + "context" + "errors" + "testing" + + "github.com/cloudwego/eino/adk" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func TestEinoCheckpointResumeHandlerSkipsWithoutCheckpoint(t *testing.T) { + called := false + handler := newEinoCheckpointResumeHandler(einoCheckpointResumeHandlerConfig{ + Resume: func(string) (*adk.AsyncIterator[*adk.AgentEvent], error) { + called = true + return nil, nil + }, + }) + if iter := handler.TryResume(); iter != nil { + t.Fatalf("iter = %#v, want nil", iter) + } + if called { + t.Fatal("resume should not be called without checkpoint state") + } +} + +func TestEinoCheckpointResumeHandlerResumesExistingCheckpoint(t *testing.T) { + store, err := newFileCheckPointStore(t.TempDir()) + if err != nil { + t.Fatal(err) + } + if err := store.Set(context.Background(), "cp-1", []byte("checkpoint")); err != nil { + t.Fatal(err) + } + var progressMessages []string + var resumedID string + core, logs := observer.New(zap.InfoLevel) + wantIter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + defer gen.Close() + handler := newEinoCheckpointResumeHandler(einoCheckpointResumeHandlerConfig{ + Context: context.Background(), + ConversationID: "conv-1", + OrchMode: "deep", + Store: store, + CheckPointID: "cp-1", + Logger: zap.New(core), + Progress: func(eventType, message string, data interface{}) { + if eventType != "progress" { + return + } + progressMessages = append(progressMessages, message) + m, _ := data.(map[string]interface{}) + if m["conversationId"] != "conv-1" || m["orchestration"] != "deep" || m["checkPointID"] != "cp-1" { + t.Fatalf("progress data = %#v", m) + } + }, + Resume: func(checkPointID string) (*adk.AsyncIterator[*adk.AgentEvent], error) { + resumedID = checkPointID + return wantIter, nil + }, + }) + + got := handler.TryResume() + if got != wantIter { + t.Fatalf("iter = %#v, want resume iterator", got) + } + if resumedID != "cp-1" { + t.Fatalf("resumed id = %q", resumedID) + } + if len(progressMessages) != 1 || progressMessages[0] != "检测到断点,正在从中断节点恢复执行..." { + t.Fatalf("progress messages = %#v", progressMessages) + } + if logs.FilterMessage("eino runner: resume from checkpoint").Len() != 1 { + t.Fatalf("expected resume log, got %d", logs.Len()) + } +} + +func TestEinoCheckpointResumeHandlerFallsBackOnResumeError(t *testing.T) { + store, err := newFileCheckPointStore(t.TempDir()) + if err != nil { + t.Fatal(err) + } + if err := store.Set(context.Background(), "cp-1", []byte("checkpoint")); err != nil { + t.Fatal(err) + } + var progressMessages []string + core, logs := observer.New(zap.WarnLevel) + handler := newEinoCheckpointResumeHandler(einoCheckpointResumeHandlerConfig{ + Context: context.Background(), + Store: store, + CheckPointID: "cp-1", + Logger: zap.New(core), + Progress: func(eventType, message string, _ interface{}) { + if eventType == "progress" { + progressMessages = append(progressMessages, message) + } + }, + Resume: func(string) (*adk.AsyncIterator[*adk.AgentEvent], error) { + return nil, errors.New("resume failed") + }, + }) + + if iter := handler.TryResume(); iter != nil { + t.Fatalf("iter = %#v, want nil fallback", iter) + } + if len(progressMessages) != 2 || progressMessages[1] != "断点恢复失败,已回退为全新执行。" { + t.Fatalf("progress messages = %#v", progressMessages) + } + if logs.FilterMessage("eino runner: resume failed, fallback to fresh run").Len() != 1 { + t.Fatalf("expected fallback log, got %d", logs.Len()) + } +} + +func TestEinoCheckpointResumeHandlerLogsPreflightError(t *testing.T) { + store, err := newFileCheckPointStore(t.TempDir()) + if err != nil { + t.Fatal(err) + } + core, logs := observer.New(zap.WarnLevel) + handler := newEinoCheckpointResumeHandler(einoCheckpointResumeHandlerConfig{ + Context: context.Background(), + Store: store, + CheckPointID: "bad/id", + Logger: zap.New(core), + Resume: func(string) (*adk.AsyncIterator[*adk.AgentEvent], error) { + t.Fatal("resume should not be called after preflight error") + return nil, nil + }, + }) + if iter := handler.TryResume(); iter != nil { + t.Fatalf("iter = %#v, want nil", iter) + } + if logs.FilterMessage("eino checkpoint preflight get failed").Len() != 1 { + t.Fatalf("expected preflight warning, got %d", logs.Len()) + } +} diff --git a/internal/multiagent/eino_checkpoint_runtime.go b/internal/multiagent/eino_checkpoint_runtime.go new file mode 100644 index 00000000..130f408d --- /dev/null +++ b/internal/multiagent/eino_checkpoint_runtime.go @@ -0,0 +1,38 @@ +package multiagent + +import ( + "path/filepath" + "strings" + + "go.uber.org/zap" +) + +type einoCheckpointRuntime struct { + Store *fileCheckPointStore + CheckPointID string +} + +func newEinoCheckpointRuntime(checkpointDir, conversationID, orchMode string, logger *zap.Logger) *einoCheckpointRuntime { + checkpointDir = strings.TrimSpace(checkpointDir) + if checkpointDir == "" { + return nil + } + cpDir := filepath.Join(checkpointDir, sanitizeEinoPathSegment(conversationID)) + store, err := newFileCheckPointStore(cpDir) + if err != nil { + if logger != nil { + logger.Warn("eino checkpoint store disabled", zap.String("dir", cpDir), zap.Error(err)) + } + return nil + } + checkPointID := buildEinoCheckpointID(orchMode) + if logger != nil { + logger.Info("eino runner: checkpoint store enabled", + zap.String("dir", cpDir), + zap.String("checkPointID", checkPointID)) + } + return &einoCheckpointRuntime{ + Store: store, + CheckPointID: checkPointID, + } +} diff --git a/internal/multiagent/eino_checkpoint_runtime_test.go b/internal/multiagent/eino_checkpoint_runtime_test.go new file mode 100644 index 00000000..23196ee2 --- /dev/null +++ b/internal/multiagent/eino_checkpoint_runtime_test.go @@ -0,0 +1,48 @@ +package multiagent + +import ( + "os" + "strings" + "testing" + + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func TestNewEinoCheckpointRuntimeDisabledWithoutDir(t *testing.T) { + if got := newEinoCheckpointRuntime(" ", "conv-1", "deep", nil); got != nil { + t.Fatalf("runtime = %#v, want nil", got) + } +} + +func TestNewEinoCheckpointRuntimeCreatesStore(t *testing.T) { + core, logs := observer.New(zap.InfoLevel) + runtime := newEinoCheckpointRuntime(t.TempDir(), "conv/1", "deep", zap.New(core)) + if runtime == nil || runtime.Store == nil { + t.Fatal("expected checkpoint runtime with store") + } + if runtime.CheckPointID != buildEinoCheckpointID("deep") { + t.Fatalf("checkpoint id = %q", runtime.CheckPointID) + } + if !strings.Contains(runtime.Store.dir, sanitizeEinoPathSegment("conv/1")) { + t.Fatalf("store dir = %q, want sanitized conversation segment", runtime.Store.dir) + } + if logs.FilterMessage("eino runner: checkpoint store enabled").Len() != 1 { + t.Fatalf("expected enabled log, got %d", logs.Len()) + } +} + +func TestNewEinoCheckpointRuntimeLogsCreateFailure(t *testing.T) { + filePath := t.TempDir() + "/not-a-dir" + if err := os.WriteFile(filePath, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + core, logs := observer.New(zap.WarnLevel) + runtime := newEinoCheckpointRuntime(filePath, "conv-1", "deep", zap.New(core)) + if runtime != nil { + t.Fatalf("runtime = %#v, want nil", runtime) + } + if logs.FilterMessage("eino checkpoint store disabled").Len() != 1 { + t.Fatalf("expected disabled log, got %d", logs.Len()) + } +} diff --git a/internal/multiagent/eino_context_overflow_retry.go b/internal/multiagent/eino_context_overflow_retry.go new file mode 100644 index 00000000..b242b440 --- /dev/null +++ b/internal/multiagent/eino_context_overflow_retry.go @@ -0,0 +1,90 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "go.uber.org/zap" +) + +type einoContextOverflowRetryConfig struct { + Context context.Context + ConversationID string + OrchMode string + Args *einoADKRunLoopArgs + BaseMsgs []adk.Message + Progress func(eventType, message string, data interface{}) + Logger *zap.Logger +} + +type einoContextOverflowRetryResult struct { + Handled bool + RestartMsgs []adk.Message + ContextSrc einoRunRestartContextSource +} + +type einoContextOverflowRetryHandler struct { + cfg einoContextOverflowRetryConfig + retried bool +} + +func newEinoContextOverflowRetryHandler(cfg einoContextOverflowRetryConfig) *einoContextOverflowRetryHandler { + if cfg.Context == nil { + cfg.Context = context.Background() + } + if cfg.Args == nil { + cfg.Args = &einoADKRunLoopArgs{} + } + return &einoContextOverflowRetryHandler{cfg: cfg} +} + +func (h *einoContextOverflowRetryHandler) Prepare( + runErr error, + accumulated []adk.Message, + baseCount int, +) einoContextOverflowRetryResult { + if h == nil || !isEinoContextOverflowError(runErr) || h.retried { + return einoContextOverflowRetryResult{} + } + h.retried = true + restartMsgs, ctxSource := einoMessagesForRunRestart(h.cfg.Args, h.cfg.BaseMsgs, accumulated, baseCount) + restartMsgs = aggressiveCompactMessagesForOverflow( + h.cfg.Context, + restartMsgs, + h.cfg.Args.MaxTotalTokens, + h.cfg.Args.ModelName, + h.cfg.Args.ToolMaxBytes, + h.cfg.OrchMode, + h.cfg.Logger, + ) + if h.cfg.Logger != nil { + h.cfg.Logger.Warn("eino context overflow, retrying with aggressive compaction", + zap.Error(runErr), + zap.String("orchestration", h.cfg.OrchMode), + zap.String("contextSource", string(ctxSource)), + ) + } + emitEinoContextOverflowRetryProgress(h.cfg.Progress, h.cfg.ConversationID, h.cfg.OrchMode, ctxSource) + return einoContextOverflowRetryResult{ + Handled: true, + RestartMsgs: restartMsgs, + ContextSrc: ctxSource, + } +} + +func emitEinoContextOverflowRetryProgress( + progress func(eventType, message string, data interface{}), + conversationID, orchMode string, + ctxSource einoRunRestartContextSource, +) bool { + if progress == nil { + return false + } + progress("eino_context_overflow_retry", "上下文超限,正在激进压缩后重试…", map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "orchestration": orchMode, + "contextSource": string(ctxSource), + }) + return true +} diff --git a/internal/multiagent/eino_context_overflow_retry_test.go b/internal/multiagent/eino_context_overflow_retry_test.go new file mode 100644 index 00000000..ca2ef603 --- /dev/null +++ b/internal/multiagent/eino_context_overflow_retry_test.go @@ -0,0 +1,90 @@ +package multiagent + +import ( + "context" + "errors" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func TestEinoContextOverflowRetryHandlerPreparesOnce(t *testing.T) { + baseMsgs := []adk.Message{ + schema.UserMessage("base"), + } + accumulated := []adk.Message{ + schema.UserMessage("base"), + schema.AssistantMessage("partial", nil), + } + var gotType, gotMessage string + var gotData map[string]interface{} + core, logs := observer.New(zap.WarnLevel) + handler := newEinoContextOverflowRetryHandler(einoContextOverflowRetryConfig{ + Context: context.Background(), + ConversationID: "conv-1", + OrchMode: "deep_agent", + Args: &einoADKRunLoopArgs{}, + BaseMsgs: baseMsgs, + Progress: func(eventType, message string, data interface{}) { + gotType = eventType + gotMessage = message + var ok bool + gotData, ok = data.(map[string]interface{}) + if !ok { + t.Fatalf("progress data type = %T, want map[string]interface{}", data) + } + }, + Logger: zap.New(core), + }) + + result := handler.Prepare(errors.New("context length exceeded"), accumulated, len(baseMsgs)) + if !result.Handled { + t.Fatal("handled = false, want true") + } + if result.ContextSrc != einoRestartContextAccumulated { + t.Fatalf("context source = %q, want %q", result.ContextSrc, einoRestartContextAccumulated) + } + if len(result.RestartMsgs) != len(accumulated) { + t.Fatalf("restart message count = %d, want %d", len(result.RestartMsgs), len(accumulated)) + } + if gotType != "eino_context_overflow_retry" { + t.Fatalf("event type = %q, want eino_context_overflow_retry", gotType) + } + if gotMessage != "上下文超限,正在激进压缩后重试…" { + t.Fatalf("message = %q", gotMessage) + } + assertContextOverflowMapValue(t, gotData, "conversationId", "conv-1") + assertContextOverflowMapValue(t, gotData, "source", "eino") + assertContextOverflowMapValue(t, gotData, "orchestration", "deep_agent") + assertContextOverflowMapValue(t, gotData, "contextSource", string(einoRestartContextAccumulated)) + if logs.FilterMessage("eino context overflow, retrying with aggressive compaction").Len() != 1 { + t.Fatalf("expected one context overflow retry log, got %d", logs.Len()) + } + + second := handler.Prepare(errors.New("maximum context length"), accumulated, len(baseMsgs)) + if second.Handled { + t.Fatalf("second result = %+v, want unhandled after first retry", second) + } +} + +func TestEinoContextOverflowRetryHandlerIgnoresOtherErrors(t *testing.T) { + handler := newEinoContextOverflowRetryHandler(einoContextOverflowRetryConfig{ + Context: context.Background(), + Args: &einoADKRunLoopArgs{}, + BaseMsgs: []adk.Message{schema.UserMessage("base")}, + }) + result := handler.Prepare(errors.New("HTTP 429 Too Many Requests"), nil, 0) + if result.Handled { + t.Fatalf("result = %+v, want unhandled", result) + } +} + +func assertContextOverflowMapValue(t *testing.T, data map[string]interface{}, key string, want interface{}) { + t.Helper() + if got := data[key]; got != want { + t.Fatalf("%s = %v, want %v", key, got, want) + } +} diff --git a/internal/multiagent/eino_empty_response.go b/internal/multiagent/eino_empty_response.go new file mode 100644 index 00000000..f0a10922 --- /dev/null +++ b/internal/multiagent/eino_empty_response.go @@ -0,0 +1,59 @@ +package multiagent + +import ( + "strings" + "time" + + "cyberstrike-ai/internal/config" +) + +const defaultEmptyResponseContinueMaxAttempts = 5 + +// IsEinoEmptyResponseResult 判断 Run 是否以「未捕获助手正文」占位结束(非真实用户可见回复)。 +func IsEinoEmptyResponseResult(result *RunResult) bool { + if result == nil { + return false + } + return isEinoEmptyResponseText(result.Response) +} + +func isEinoEmptyResponseText(s string) bool { + s = strings.TrimSpace(s) + if s == "" { + return false + } + return strings.Contains(s, "no assistant text was captured") || + strings.Contains(s, "未捕获到助手文本输出") +} + +// HasEinoResumeTrace 轨迹非空,续跑才有上下文可恢复。 +func HasEinoResumeTrace(result *RunResult) bool { + if result == nil { + return false + } + s := strings.TrimSpace(result.LastAgentTraceInput) + return s != "" && s != "[]" && s != "null" +} + +// EmptyResponseContinueMaxAttemptsFromConfig 无助手正文时 Handler 层退避续跑上限;0=默认 5。 +func EmptyResponseContinueMaxAttemptsFromConfig(mw *config.MultiAgentEinoMiddlewareConfig) int { + if mw != nil && mw.EmptyResponseContinueMaxAttempts > 0 { + return mw.EmptyResponseContinueMaxAttempts + } + return defaultEmptyResponseContinueMaxAttempts +} + +// EmptyResponseContinueBackoff 与 run_retry 相同指数退避(2s, 4s, 8s… capped)。 +func EmptyResponseContinueBackoff(attempt int, mw *config.MultiAgentEinoMiddlewareConfig) time.Duration { + maxBackoff := defaultEinoRunRetryMaxBackoff + if mw != nil && mw.RunRetryMaxBackoffSec > 0 { + maxBackoff = time.Duration(mw.RunRetryMaxBackoffSec) * time.Second + } + return einoTransientRetryBackoff(attempt, maxBackoff) +} + +// FormatEmptyResponseContinueUserMessage 系统自动续跑时注入的 user 轮次(不写入 messages 表气泡)。 +func FormatEmptyResponseContinueUserMessage() string { + return strings.TrimSpace(`【系统自动续跑 / Auto resume】 +上一轮 Eino 会话未产出可见助手正文(可能流式中断或仅完成工具调用)。请基于已有轨迹与工具结果继续推进,并给出阶段性总结;勿重复已完成步骤。`) +} diff --git a/internal/multiagent/eino_empty_response_test.go b/internal/multiagent/eino_empty_response_test.go new file mode 100644 index 00000000..68354263 --- /dev/null +++ b/internal/multiagent/eino_empty_response_test.go @@ -0,0 +1,38 @@ +package multiagent + +import "testing" + +func TestIsEinoEmptyResponseResult(t *testing.T) { + empty := &RunResult{ + Response: "(Eino ADK single-agent session completed but no assistant text was captured. Check process details or logs.) " + + "(Eino ADK 单代理会话已完成,但未捕获到助手文本输出。请查看过程详情或日志。)", + } + if !IsEinoEmptyResponseResult(empty) { + t.Fatal("expected empty placeholder response") + } + ok := &RunResult{Response: "扫描完成,发现 2 个开放端口。"} + if IsEinoEmptyResponseResult(ok) { + t.Fatalf("expected real response, got placeholder match") + } + if IsEinoEmptyResponseResult(nil) { + t.Fatal("nil result should be false") + } +} + +func TestHasEinoResumeTrace(t *testing.T) { + if HasEinoResumeTrace(nil) { + t.Fatal("nil") + } + if HasEinoResumeTrace(&RunResult{LastAgentTraceInput: "[]"}) { + t.Fatal("enable resume on empty trace") + } + if !HasEinoResumeTrace(&RunResult{LastAgentTraceInput: `[{"role":"user","content":"hi"}]`}) { + t.Fatal("expected resume trace") + } +} + +func TestEmptyResponseContinueMaxAttemptsFromConfig(t *testing.T) { + if got := EmptyResponseContinueMaxAttemptsFromConfig(nil); got != defaultEmptyResponseContinueMaxAttempts { + t.Fatalf("default: got %d want %d", got, defaultEmptyResponseContinueMaxAttempts) + } +} diff --git a/internal/multiagent/eino_execute_failure_format_test.go b/internal/multiagent/eino_execute_failure_format_test.go new file mode 100644 index 00000000..5b831d81 --- /dev/null +++ b/internal/multiagent/eino_execute_failure_format_test.go @@ -0,0 +1,152 @@ +package multiagent + +import ( + "context" + "errors" + "io" + "strings" + "testing" + + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/security" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/schema" +) + +type mockStreamingShellExitFail struct { + output string + code int +} + +func (m *mockStreamingShellExitFail) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4) + go func() { + defer outW.Close() + if m.output != "" { + _ = outW.Send(&filesystem.ExecuteResponse{Output: m.output}, nil) + } + code := m.code + _ = outW.Send(&filesystem.ExecuteResponse{ExitCode: &code}, nil) + }() + return outR, nil +} + +func TestEinoStreamingShellWrap_CommandFailureFormat(t *testing.T) { + inner := &mockStreamingShellExitFail{ + output: "sudo: a password is required\n", + code: 1, + } + notify := einomcp.NewToolInvokeNotifyHolder() + var firedBody string + var firedSuccess bool + var firedErr error + notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) { + firedBody = content + firedSuccess = success + firedErr = invokeErr + }) + wrap := &einoStreamingShellWrap{inner: inner, invokeNotify: notify} + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + var stream strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil { + stream.WriteString(resp.Output) + } + } + + if firedSuccess { + t.Fatal("expected success=false") + } + var exitErr *ExecuteExitError + if !errors.As(firedErr, &exitErr) || exitErr.Code != 1 { + t.Fatalf("expected ExecuteExitError code 1, got %v", firedErr) + } + if !strings.HasPrefix(firedBody, einomcp.ToolErrorPrefix) { + t.Fatalf("missing tool error prefix: %q", firedBody) + } + body := strings.TrimPrefix(firedBody, einomcp.ToolErrorPrefix) + if body != security.FormatCommandFailureResult(1, "sudo: a password is required\n") { + t.Fatalf("fire body = %q", body) + } + if !strings.Contains(stream.String(), "sudo:") { + t.Fatalf("stream missing sudo output: %q", stream.String()) + } + if strings.Contains(stream.String(), "command exited with non-zero") { + t.Fatalf("stream has legacy noise: %q", stream.String()) + } + if strings.Contains(stream.String(), "执行未正常结束") { + t.Fatalf("stream has abnormal tail: %q", stream.String()) + } + if !security.IsCommandFailureResult(stream.String()) { + t.Fatalf("stream missing failure status line: %q", stream.String()) + } + if tail := friendlyEinoExecuteInvokeTail(firedErr); tail != "" { + t.Fatalf("unexpected invoke tail: %q", tail) + } + if !einoToolResultIsError("execute", firedBody) { + t.Fatal("expected isError for execute failure") + } +} + +func TestFriendlyEinoExecuteInvokeTail(t *testing.T) { + if friendlyEinoExecuteInvokeTail(&ExecuteExitError{Code: 1}) != "" { + t.Fatal("exit error should not get abnormal tail") + } + if !strings.Contains(friendlyEinoExecuteInvokeTail(context.DeadlineExceeded), "Timed out") { + t.Fatal("deadline should get timeout hint") + } + if friendlyEinoExecuteInvokeTail(errors.New("broken pipe")) == "" { + t.Fatal("unexpected error should get tail") + } +} + +func TestMCPBackgroundWaitResultIsDisplayRunning(t *testing.T) { + body := `工具已提交到后台执行,但本次等待已到达上限。 + +execution_id: 3eaaa391-050b-4be1-a870-48a855923cb7 +tool: exec +status: running +wait_timeout: 10s +elapsed: 10s + +你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。` + modelFacing := einomcp.ToolErrorPrefix + body + if !einoToolResultIsError("exec", modelFacing) { + t.Fatal("soft wait timeout must remain model-facing tool error") + } + if !isMCPBackgroundWaitResult(einoToolResultBody(modelFacing)) { + t.Fatal("soft wait timeout should display as background running") + } + if got := mcpExecutionIDFromWaitResult(einoToolResultBody(modelFacing)); got != "3eaaa391-050b-4be1-a870-48a855923cb7" { + t.Fatalf("execution id = %q", got) + } + if isMCPBackgroundWaitResult("execution_id: abc\nstatus: failed\nerror: boom") { + t.Fatal("real failures must not display as background running") + } + jsonBody := `{ + "execution_id": "e98baefc-72eb-4a7e-9091-9be179a75d71", + "tool": "exec", + "status": "running" +} + +本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。` + if !isMCPBackgroundWaitResult(jsonBody) { + t.Fatal("json wait_tool_execution timeout should display as background running") + } + if got := mcpExecutionIDFromWaitResult(jsonBody); got != "e98baefc-72eb-4a7e-9091-9be179a75d71" { + t.Fatalf("json execution id = %q", got) + } +} diff --git a/internal/multiagent/eino_execute_monitor.go b/internal/multiagent/eino_execute_monitor.go new file mode 100644 index 00000000..6630b94b --- /dev/null +++ b/internal/multiagent/eino_execute_monitor.go @@ -0,0 +1,68 @@ +package multiagent + +import ( + "context" + "fmt" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/einomcp" +) + +// newEinoExecuteMonitorCallbacks 在 Eino filesystem execute 开始/结束时写入 MCP 监控库并 recorder(executionId), +// 与 CallTool 路径一致,使监控页能展示「执行中」状态。 +func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) ( + begin func(toolCallID, command string) string, + appendPartial func(executionID, toolCallID, chunk string), + registerCancel func(executionID string, cancel context.CancelFunc), + unregisterCancel func(executionID string), + finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error), +) { + begin = func(toolCallID, command string) string { + if ag == nil { + return "" + } + args := map[string]interface{}{"command": command} + id := ag.BeginLocalToolExecution(ctx, "execute", args) + if id != "" && recorder != nil { + recorder(id, toolCallID) + } + return id + } + appendPartial = func(executionID, toolCallID, chunk string) { + if ag == nil || executionID == "" || chunk == "" { + return + } + ag.AppendLocalToolExecutionPartialOutput(executionID, chunk) + } + registerCancel = func(executionID string, cancel context.CancelFunc) { + if ag == nil || executionID == "" || cancel == nil { + return + } + ag.RegisterLocalToolExecutionCancel(executionID, cancel) + } + unregisterCancel = func(executionID string) { + if ag == nil || executionID == "" { + return + } + ag.UnregisterLocalToolExecutionCancel(executionID) + } + finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) { + if ag == nil { + return + } + var err error + if !success { + if invokeErr != nil { + err = invokeErr + } else { + err = fmt.Errorf("execute failed") + } + } + args := map[string]interface{}{"command": command} + id := ag.FinishLocalToolExecution(ctx, executionID, "execute", args, stdout, err) + if id != "" && recorder != nil && executionID == "" { + recorder(id, toolCallID) + } + } + return begin, appendPartial, registerCancel, unregisterCancel, finish +} diff --git a/internal/multiagent/eino_execute_stdout_suppressor.go b/internal/multiagent/eino_execute_stdout_suppressor.go new file mode 100644 index 00000000..455ce409 --- /dev/null +++ b/internal/multiagent/eino_execute_stdout_suppressor.go @@ -0,0 +1,57 @@ +package multiagent + +import ( + "strings" + "sync" +) + +type einoExecuteStdoutSuppressor struct { + mu sync.Mutex + pending string +} + +func newEinoExecuteStdoutSuppressor() *einoExecuteStdoutSuppressor { + return &einoExecuteStdoutSuppressor{} +} + +func (s *einoExecuteStdoutSuppressor) Record(toolName, stdout string, isErr bool) { + if s == nil || isErr || !strings.EqualFold(strings.TrimSpace(toolName), "execute") { + return + } + t := strings.TrimSpace(stdout) + if t == "" { + return + } + s.mu.Lock() + s.pending = t + s.mu.Unlock() +} + +func (s *einoExecuteStdoutSuppressor) Peek() string { + if s == nil { + return "" + } + s.mu.Lock() + defer s.mu.Unlock() + return s.pending +} + +func (s *einoExecuteStdoutSuppressor) Consume() string { + if s == nil { + return "" + } + s.mu.Lock() + defer s.mu.Unlock() + out := s.pending + s.pending = "" + return out +} + +func (s *einoExecuteStdoutSuppressor) Clear() { + if s == nil { + return + } + s.mu.Lock() + s.pending = "" + s.mu.Unlock() +} diff --git a/internal/multiagent/eino_execute_stdout_suppressor_test.go b/internal/multiagent/eino_execute_stdout_suppressor_test.go new file mode 100644 index 00000000..f4e99e01 --- /dev/null +++ b/internal/multiagent/eino_execute_stdout_suppressor_test.go @@ -0,0 +1,42 @@ +package multiagent + +import "testing" + +func TestEinoExecuteStdoutSuppressorRecordsOnlySuccessfulExecute(t *testing.T) { + s := newEinoExecuteStdoutSuppressor() + s.Record("read_file", "file body", false) + if got := s.Peek(); got != "" { + t.Fatalf("non-execute should not be recorded, got %q", got) + } + s.Record("execute", "failed", true) + if got := s.Peek(); got != "" { + t.Fatalf("failed execute should not be recorded, got %q", got) + } + s.Record(" execute ", " hello\n", false) + if got := s.Peek(); got != "hello" { + t.Fatalf("Peek = %q, want hello", got) + } +} + +func TestEinoExecuteStdoutSuppressorConsumeAndClear(t *testing.T) { + s := newEinoExecuteStdoutSuppressor() + s.Record("execute", "stdout", false) + if got := s.Peek(); got != "stdout" { + t.Fatalf("Peek = %q, want stdout", got) + } + if got := s.Peek(); got != "stdout" { + t.Fatalf("Peek should not clear, got %q", got) + } + if got := s.Consume(); got != "stdout" { + t.Fatalf("Consume = %q, want stdout", got) + } + if got := s.Peek(); got != "" { + t.Fatalf("Consume should clear, got %q", got) + } + + s.Record("execute", "again", false) + s.Clear() + if got := s.Consume(); got != "" { + t.Fatalf("Clear should remove pending value, got %q", got) + } +} diff --git a/internal/multiagent/eino_execute_streaming_wrap.go b/internal/multiagent/eino_execute_streaming_wrap.go new file mode 100644 index 00000000..0e7af38b --- /dev/null +++ b/internal/multiagent/eino_execute_streaming_wrap.go @@ -0,0 +1,409 @@ +package multiagent + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// prependPythonUnbufferedEnv 为 /bin/sh -c 注入 PYTHONUNBUFFERED=1。 +// eino-ext local 对流式 stdout 使用 bufio 按「行」推送;python3 写管道时默认块缓冲,print 长期留在用户态缓冲, +// 管道里收不到换行,表现为长时间无输出直至超时或退出。若命令里已出现 PYTHONUNBUFFERED 则不再覆盖。 +func prependPythonUnbufferedEnv(shellCommand string) string { + if strings.TrimSpace(shellCommand) == "" { + return shellCommand + } + if strings.Contains(strings.ToUpper(shellCommand), "PYTHONUNBUFFERED") { + return shellCommand + } + return "export PYTHONUNBUFFERED=1\n" + shellCommand +} + +// einoExecuteTimeoutUserHint 与写入 ADK 工具消息(模型可见)及 SSE tool_result 尾标一致。 +func einoExecuteTimeoutUserHint() string { + return "已超时终止 · Timed out" +} + +// einoExecuteRecvErrIsToolTimeout 判断 Recv 错误是否由 agent.tool_timeout_minutes 触发。 +// WithTimeout 到期后 local 侧常报 canceled / exit -1,但 execCtx.Err() 仍为 DeadlineExceeded。 +func einoExecuteRecvErrIsToolTimeout(rerr error, tctx context.Context) bool { + if tctx != nil && errors.Is(tctx.Err(), context.DeadlineExceeded) { + return true + } + return errors.Is(rerr, context.DeadlineExceeded) +} + +// einoStreamingShellWrap 包装 Eino filesystem 使用的 StreamingShell(cloudwego eino-ext local.Local)。 +// 官方 execute 工具默认走 ExecuteStreaming 且不设 RunInBackendGround;末尾带 & 时子进程仍与管道相连, +// streamStdout 按行读取会在无换行输出时长时间阻塞(与 MCP 工具 exec 的独立实现不同)。 +// 对「完全后台」命令自动开启 RunInBackendGround,与 local.runCmdInBackground 行为对齐。 +// +// 使用 Pipe 将内层流转发给调用方:在 inner EOF 后、关闭 Pipe 前同步调用 ToolInvokeNotify.Fire, +// run loop 收到 Fire 后立即推送 tool_result(toolResultSent 去重),避免 ADK Tool 事件迟到时 UI 卡在「执行中」。 +// +// 若 inner 在校验阶段直接返回 error(未建立 reader),不会进入下方 goroutine,也必须 Fire; +// 否则 pending tool_call 要等整轮 run 结束才被 force-close,与已展示的助手/工具软错误文案不同步。 +type einoStreamingShellWrap struct { + inner filesystem.StreamingShell + invokeNotify *einomcp.ToolInvokeNotifyHolder + einoAgentName string + // outputChunk 可选;非 nil 时在收到内层 ExecuteResponse 片段时推送,与 MCP 工具的 tool_result_delta 一致(需有效 toolCallId)。 + outputChunk func(toolName, toolCallID, chunk string) + // toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。 + toolTimeoutMinutes int + // toolWaitTimeoutSeconds 与 agent.tool_wait_timeout_seconds 对齐;>0 时本轮等待到期后返回 execution_id,shell 继续后台运行。 + toolWaitTimeoutSeconds int + // shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。 + shellNoOutputTimeoutSec int + // beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。 + beginMonitor func(toolCallID, command string) string + appendPartialMonitor func(executionID, toolCallID, chunk string) + registerCancelMonitor func(executionID string, cancel context.CancelFunc) + unregisterCancelMonitor func(executionID string) + finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) +} + +func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + if w.inner == nil { + return nil, fmt.Errorf("einoStreamingShellWrap: inner shell is nil") + } + if input == nil { + return w.inner.ExecuteStreaming(ctx, nil) + } + req := *input + userCmd := strings.TrimSpace(req.Command) + tid := strings.TrimSpace(compose.GetToolCallID(ctx)) + agentTag := strings.TrimSpace(w.einoAgentName) + if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround { + req.RunInBackendGround = true + } + req.Command = prependPythonUnbufferedEnv(req.Command) + convID := mcp.MCPConversationIDFromContext(ctx) + execReg := mcp.EinoExecuteRunRegistryFromContext(ctx) + + var monitorExecID string + if w.beginMonitor != nil { + monitorExecID = w.beginMonitor(tid, userCmd) + } + if monitorExecID != "" && convID != "" { + if toolReg := mcp.ToolRunRegistryFromContext(ctx); toolReg != nil { + toolReg.RegisterRunningTool(convID, monitorExecID) + } + } + toolRunReg := mcp.ToolRunRegistryFromContext(ctx) + + execCtx, execCancel := context.WithCancel(ctx) + var timeoutCancel context.CancelFunc + if w.toolTimeoutMinutes > 0 { + execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute) + } + if monitorExecID != "" && w.registerCancelMonitor != nil { + w.registerCancelMonitor(monitorExecID, execCancel) + } + if execReg != nil && convID != "" { + execReg.RegisterActiveEinoExecute(convID, execCancel) + } + + sr, err := w.inner.ExecuteStreaming(execCtx, &req) + if err != nil { + if timeoutCancel != nil { + timeoutCancel() + } + if execCancel != nil { + execCancel() + } + if monitorExecID != "" && w.unregisterCancelMonitor != nil { + w.unregisterCancelMonitor(monitorExecID) + } + if einoExecuteRecvErrIsToolTimeout(err, execCtx) { + hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n" + if w.finishMonitor != nil { + w.finishMonitor(monitorExecID, tid, userCmd, hint, false, context.DeadlineExceeded) + } + if w.invokeNotify != nil && tid != "" { + w.invokeNotify.Fire(tid, "execute", agentTag, false, hint, context.DeadlineExceeded) + } + return schema.StreamReaderFromArray([]*filesystem.ExecuteResponse{{Output: hint}}), nil + } + if w.finishMonitor != nil { + w.finishMonitor(monitorExecID, tid, userCmd, "", false, err) + } + if w.invokeNotify != nil && tid != "" { + w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err) + } + return nil, err + } + if sr == nil { + if timeoutCancel != nil { + timeoutCancel() + } + if execCancel != nil { + execCancel() + } + return sr, nil + } + + outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32) + + go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int, waitTimeoutSec int) { + var innerCloseOnce sync.Once + closeInner := func() { + innerCloseOnce.Do(func() { inner.Close() }) + } + defer closeInner() + if timeoutCleanup != nil { + defer timeoutCleanup() + } + if cancel != nil { + defer cancel() + } + if reg != nil && conversationID != "" { + defer reg.UnregisterActiveEinoExecute(conversationID) + } + if toolReg != nil && conversationID != "" && execID != "" { + defer toolReg.UnregisterRunningTool(conversationID, execID) + } + if w.unregisterCancelMonitor != nil && execID != "" { + defer w.unregisterCancelMonitor(execID) + } + + // ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。 + stopWatch := make(chan struct{}) + go func() { + select { + case <-tctx.Done(): + closeInner() + case <-stopWatch: + } + }() + defer close(stopWatch) + + var sb strings.Builder + success := true + var invokeErr error + exitCode := 0 + hasExitCode := false + softReturned := false + var outCloseOnce sync.Once + closeOut := func() { + outCloseOnce.Do(func() { outW.Close() }) + } + defer closeOut() + sendOut := func(resp *filesystem.ExecuteResponse, err error) bool { + if softReturned { + return false + } + return outW.Send(resp, err) + } + + idleWatch := security.NewShellInactivityWatch(noOutputSec) + if idleWatch != nil { + defer idleWatch.Stop() + } + var waitTimeoutCh <-chan time.Time + var waitTimer *time.Timer + if waitTimeoutSec > 0 { + waitTimer = time.NewTimer(time.Duration(waitTimeoutSec) * time.Second) + waitTimeoutCh = waitTimer.C + defer waitTimer.Stop() + } + + type execRecvMsg struct { + resp *filesystem.ExecuteResponse + err error + } + recvCh := make(chan execRecvMsg, 1) + go func() { + for { + resp, rerr := inner.Recv() + recvCh <- execRecvMsg{resp: resp, err: rerr} + if rerr != nil { + return + } + } + }() + + fireInactivityTimeout := func() { + success = false + invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) + msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec) + _ = sendOut(&filesystem.ExecuteResponse{Output: msg}, nil) + sb.WriteString(msg) + if w.appendPartialMonitor != nil && execID != "" { + w.appendPartialMonitor(execID, toolCallID, msg) + } + if w.outputChunk != nil && toolCallID != "" { + w.outputChunk("execute", toolCallID, msg) + } + if cancel != nil { + cancel() + } + closeInner() + } + + recvLoop: + for { + var idleCh <-chan struct{} + if idleWatch != nil { + idleCh = idleWatch.Expired + } + select { + case <-idleCh: + fireInactivityTimeout() + break recvLoop + case <-waitTimeoutCh: + if execID != "" && !softReturned { + msg := einoExecuteSoftWaitTimeoutResult(execID, waitTimeoutSec) + _ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil) + softReturned = true + closeOut() + } + waitTimeoutCh = nil + case msg := <-recvCh: + rerr := msg.err + resp := msg.resp + if errors.Is(rerr, io.EOF) { + break recvLoop + } + if rerr != nil { + success = false + invokeErr = rerr + if einoExecuteRecvErrIsToolTimeout(rerr, tctx) { + invokeErr = context.DeadlineExceeded + break recvLoop + } + if errors.Is(rerr, context.Canceled) || (tctx != nil && errors.Is(tctx.Err(), context.Canceled)) { + invokeErr = context.Canceled + break recvLoop + } + _ = sendOut(nil, rerr) + break recvLoop + } + if resp != nil { + if resp.ExitCode != nil { + hasExitCode = true + exitCode = *resp.ExitCode + continue + } + var appended string + if resp.Output != "" { + if security.IsLegacyShellExitNoise(resp.Output) { + continue + } + if idleWatch != nil { + idleWatch.Bump() + } + sb.WriteString(resp.Output) + appended = resp.Output + if w.appendPartialMonitor != nil && execID != "" { + w.appendPartialMonitor(execID, toolCallID, appended) + } + } + if w.outputChunk != nil && strings.TrimSpace(appended) != "" { + w.outputChunk("execute", toolCallID, appended) + } + if sendOut(resp, nil) { + success = false + invokeErr = fmt.Errorf("execute stream closed by consumer") + break recvLoop + } + } + } + } + + if success && hasExitCode && exitCode != 0 { + success = false + invokeErr = &ExecuteExitError{Code: exitCode} + } + // WithTimeout 触发后,子进程常被信号结束,local 侧多报 exit -1 / canceled,错误链里不一定带 DeadlineExceeded。 + // 用执行所用 ctx 归一化,便于 UI 展示「超时」而非含糊的 -1。 + if tctx != nil && errors.Is(tctx.Err(), context.DeadlineExceeded) { + success = false + invokeErr = context.DeadlineExceeded + } + // 用户「中断并继续」终止 execute:合并说明进工具结果(与 MCP CancelToolExecutionWithNote 一致)。 + partialStreamed := sb.String() + var abortNote string + if reg != nil && conversationID != "" && (invokeErr != nil || errors.Is(tctx.Err(), context.Canceled)) { + if note := reg.TakeEinoExecuteAbortNote(conversationID); note != "" { + abortNote = note + merged := mcp.MergePartialToolOutputAndAbortNote(partialStreamed, note) + sb.Reset() + sb.WriteString(merged) + if invokeErr == nil { + success = false + invokeErr = context.Canceled + } + } + } + // ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。 + if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) { + hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n" + _ = sendOut(&filesystem.ExecuteResponse{Output: hint}, nil) + if w.appendPartialMonitor != nil && execID != "" { + w.appendPartialMonitor(execID, toolCallID, hint) + } + if w.outputChunk != nil && tid != "" { + w.outputChunk("execute", tid, hint) + } + sb.WriteString(hint) + } + // 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。 + if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" { + if partialStreamed != "" { + _ = sendOut(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil) + } else if text := strings.TrimSpace(sb.String()); text != "" { + _ = sendOut(&filesystem.ExecuteResponse{Output: text + "\n"}, nil) + } + } + rawOutput := sb.String() + fireBody := rawOutput + if !success && hasExitCode && exitCode != 0 { + statusLine := security.ExecuteFailureStatusLine(exitCode) + if !strings.Contains(rawOutput, "命令执行失败:") { + _ = sendOut(&filesystem.ExecuteResponse{Output: statusLine}, nil) + if w.appendPartialMonitor != nil && execID != "" { + w.appendPartialMonitor(execID, toolCallID, statusLine) + } + sb.WriteString(statusLine) + } + fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput) + } + if w.finishMonitor != nil { + w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr) + } + if w.invokeNotify != nil { + if !softReturned { + w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr) + } + } + }(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec, w.toolWaitTimeoutSeconds) + + return outR, nil +} + +func einoExecuteSoftWaitTimeoutResult(executionID string, waitTimeoutSec int) string { + waitText := "configured wait timeout" + if waitTimeoutSec > 0 { + waitText = fmt.Sprintf("%ds", waitTimeoutSec) + } + return fmt.Sprintf(`工具已提交到后台执行,当前仍在运行。 + +execution_id: %s +status: running +wait_timeout: %s + +你可以继续推理、改用其他工具,或调用 get_tool_execution / wait_tool_execution 读取 partial_output 并继续等待;也可以调用 cancel_tool_execution 取消。`, executionID, waitText) +} diff --git a/internal/multiagent/eino_execute_streaming_wrap_test.go b/internal/multiagent/eino_execute_streaming_wrap_test.go new file mode 100644 index 00000000..54549b0c --- /dev/null +++ b/internal/multiagent/eino_execute_streaming_wrap_test.go @@ -0,0 +1,431 @@ +package multiagent + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/mcp" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/schema" +) + +type mockStreamingShell struct { + immediateErr error + recvErr error + output string + called bool + lastCommand string +} + +func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + m.called = true + if input != nil { + m.lastCommand = input.Command + } + if m.immediateErr != nil { + return nil, m.immediateErr + } + outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4) + go func() { + defer outW.Close() + if strings.TrimSpace(m.output) != "" { + _ = outW.Send(&filesystem.ExecuteResponse{Output: m.output}, nil) + } + if m.recvErr != nil { + _ = outW.Send(nil, m.recvErr) + } + }() + return outR, nil +} + +func TestEinoStreamingShellWrap_PreparesNonInteractiveCommand(t *testing.T) { + inner := &mockStreamingShell{output: "ok\n"} + wrap := &einoStreamingShellWrap{inner: inner} + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "echo ok"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + for { + _, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + } + if !strings.Contains(inner.lastCommand, "PYTHONUNBUFFERED=1") { + t.Fatalf("missing python unbuffer in inner command: %q", inner.lastCommand) + } +} + +func TestEinoStreamingShellWrap_NoOutputTimeout(t *testing.T) { + inner := &mockStreamingShellHanging{} + notify := einomcp.NewToolInvokeNotifyHolder() + var fired string + notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) { + fired = content + }) + wrap := &einoStreamingShellWrap{ + inner: inner, + invokeNotify: notify, + shellNoOutputTimeoutSec: 1, + } + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil { + got.WriteString(resp.Output) + } + } + if !inner.called { + t.Fatal("inner shell should run (no command blacklist)") + } + out := got.String() + if !strings.Contains(out, "没有新的输出") && !strings.Contains(out, "no new output") { + t.Fatalf("expected inactivity timeout message, got: %q notify=%q", out, fired) + } +} + +type mockStreamingShellPartialThenHang struct { + called bool +} + +func (m *mockStreamingShellPartialThenHang) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + m.called = true + outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4) + go func() { + _ = outW.Send(&filesystem.ExecuteResponse{Output: "[sudo] password:\n"}, nil) + <-ctx.Done() + outW.Close() + }() + return outR, nil +} + +func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) { + inner := &mockStreamingShellPartialThenHang{} + wrap := &einoStreamingShellWrap{ + inner: inner, + shellNoOutputTimeoutSec: 1, + } + start := time.Now() + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil { + got.WriteString(resp.Output) + } + } + if time.Since(start) > 5*time.Second { + t.Fatalf("expected inactivity timeout ~1s, took %v", time.Since(start)) + } + if !strings.Contains(got.String(), "没有新的输出") && !strings.Contains(got.String(), "no new output") { + t.Fatalf("expected inactivity message, got: %q", got.String()) + } +} + +func TestEinoStreamingShellWrap_SoftWaitTimeoutReturnsExecutionIDAndKeepsRunning(t *testing.T) { + inner := &mockStreamingShellPartialThenHang{} + partialCh := make(chan string, 4) + cancelCh := make(chan context.CancelFunc, 1) + unregistered := make(chan string, 1) + wrap := &einoStreamingShellWrap{ + inner: inner, + toolWaitTimeoutSeconds: 1, + beginMonitor: func(toolCallID, command string) string { + return "exec-soft-wait" + }, + appendPartialMonitor: func(executionID, toolCallID, chunk string) { + partialCh <- chunk + }, + registerCancelMonitor: func(executionID string, cancel context.CancelFunc) { + if executionID == "exec-soft-wait" { + cancelCh <- cancel + } + }, + unregisterCancelMonitor: func(executionID string) { + unregistered <- executionID + }, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "sudo whoami"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil { + got.WriteString(resp.Output) + } + } + body := got.String() + if !strings.Contains(body, "execution_id: exec-soft-wait") || !strings.Contains(body, "status: running") { + t.Fatalf("expected background execution marker, got: %q", body) + } + select { + case chunk := <-partialCh: + if !strings.Contains(chunk, "[sudo] password") { + t.Fatalf("unexpected partial chunk: %q", chunk) + } + default: + t.Fatal("expected streamed partial output before soft wait return") + } + if !inner.called { + t.Fatal("inner shell did not run") + } + select { + case registeredCancel := <-cancelCh: + registeredCancel() + case <-time.After(time.Second): + t.Fatal("expected execution cancel registration") + } + select { + case id := <-unregistered: + if id != "exec-soft-wait" { + t.Fatalf("unexpected unregistered id: %q", id) + } + case <-time.After(time.Second): + t.Fatal("expected execution cancel unregister") + } +} + +type mockStreamingShellHanging struct { + called bool +} + +func (m *mockStreamingShellHanging) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + m.called = true + outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4) + go func() { + <-ctx.Done() + outW.Close() + }() + return outR, nil +} + +func TestEinoExecuteRecvErrIsToolTimeout(t *testing.T) { + tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) + defer cancel() + time.Sleep(2 * time.Millisecond) + <-tctx.Done() + + if !einoExecuteRecvErrIsToolTimeout(context.Canceled, tctx) { + t.Fatal("expected canceled recv with deadline exec ctx to count as tool timeout") + } + if !einoExecuteRecvErrIsToolTimeout(context.DeadlineExceeded, nil) { + t.Fatal("expected DeadlineExceeded recv without tctx") + } + if einoExecuteRecvErrIsToolTimeout(errors.New("exit status 1"), context.Background()) { + t.Fatal("unexpected timeout for generic error") + } +} + +func TestEinoStreamingShellWrap_ToolTimeoutImmediateErrIsSoft(t *testing.T) { + inner := &mockStreamingShell{immediateErr: context.DeadlineExceeded} + wrap := &einoStreamingShellWrap{ + inner: inner, + toolTimeoutMinutes: 60, + } + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"}) + if err != nil { + t.Fatalf("immediate tool timeout must return soft stream, got err: %v", err) + } + defer sr.Close() + + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("outer stream must not hard-fail, got: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) { + t.Fatalf("expected timeout hint, got: %q", got.String()) + } +} + +func TestEinoStreamingShellWrap_ToolTimeoutRecvErrIsSoft(t *testing.T) { + inner := &mockStreamingShell{recvErr: context.DeadlineExceeded} + notify := einomcp.NewToolInvokeNotifyHolder() + wrap := &einoStreamingShellWrap{ + inner: inner, + invokeNotify: notify, + toolTimeoutMinutes: 60, + } + // 生产路径由 Eino compose 注入 toolCallID;单测通过已过期 execCtx 识别 tool_timeout 软错误。 + tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) + defer cancel() + time.Sleep(2 * time.Millisecond) + <-tctx.Done() + + sr, err := wrap.ExecuteStreaming(tctx, &filesystem.ExecuteRequest{Command: "sleep 999"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("outer stream must not hard-fail on tool timeout, got: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) { + t.Fatalf("expected timeout hint in stream, got: %q", got.String()) + } +} + +func TestEinoStreamingShellWrap_CapturesOutputWithToolTimeout(t *testing.T) { + inner := &mockStreamingShell{output: "100\n"} + notify := einomcp.NewToolInvokeNotifyHolder() + var firedContent string + notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) { + firedContent = content + }) + wrap := &einoStreamingShellWrap{ + inner: inner, + invokeNotify: notify, + toolTimeoutMinutes: 60, + } + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "echo 100"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("unexpected stream error: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if !strings.Contains(got.String(), "100") { + t.Fatalf("stream output = %q, want contains 100", got.String()) + } + if !strings.Contains(firedContent, "100") { + t.Fatalf("notify content = %q, want contains 100", firedContent) + } +} + +func TestEinoStreamingShellWrap_AbortNoteDoesNotDuplicateStreamedOutput(t *testing.T) { + inner := &mockStreamingShell{output: "line1\nline2\n", recvErr: context.Canceled} + notify := einomcp.NewToolInvokeNotifyHolder() + wrap := &einoStreamingShellWrap{ + inner: inner, + invokeNotify: notify, + } + reg := &abortNoteTestRegistry{note: "改成20次"} + ctx := mcp.WithEinoExecuteRunRegistry( + mcp.WithMCPConversationID(context.Background(), "conv-abort-dup"), + reg, + ) + sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "ping -c 10 baidu.com"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("unexpected stream error: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + out := got.String() + if strings.Count(out, "line1") != 1 || strings.Count(out, "line2") != 1 { + t.Fatalf("stream duplicated stdout: %q", out) + } + if !strings.Contains(out, "改成20次") { + t.Fatalf("stream missing abort note: %q", out) + } +} + +type abortNoteTestRegistry struct { + note string +} + +func (r *abortNoteTestRegistry) RegisterActiveEinoExecute(string, context.CancelFunc) {} +func (r *abortNoteTestRegistry) UnregisterActiveEinoExecute(string) {} +func (r *abortNoteTestRegistry) AbortActiveEinoExecute(string, string) bool { return false } +func (r *abortNoteTestRegistry) TakeEinoExecuteAbortNote(string) string { return r.note } + +func TestEinoStreamingShellWrap_NonTimeoutRecvErrStillHard(t *testing.T) { + inner := &mockStreamingShell{recvErr: errors.New("broken pipe")} + wrap := &einoStreamingShellWrap{inner: inner} + sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + _, rerr := sr.Recv() + if rerr == nil || errors.Is(rerr, io.EOF) { + t.Fatal("expected hard stream error for non-timeout failure") + } +} diff --git a/internal/multiagent/eino_exit_fallback_test.go b/internal/multiagent/eino_exit_fallback_test.go new file mode 100644 index 00000000..57bba91d --- /dev/null +++ b/internal/multiagent/eino_exit_fallback_test.go @@ -0,0 +1,62 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/schema" +) + +func TestEinoExtractFallbackAssistantFromMsgs_exitToolMessage(t *testing.T) { + u := schema.UserMessage("hi") + tm := schema.ToolMessage("answer for user", "call-exit-1") + tm.ToolName = "exit" + if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{u, tm}); got != "answer for user" { + t.Fatalf("got %q", got) + } +} + +func TestEinoExtractFallbackAssistantFromMsgs_lastExitWins(t *testing.T) { + msgs := []*schema.Message{ + schema.UserMessage("hi"), + toolExitMsg("first", "c1"), + toolExitMsg("second", "c2"), + } + if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "second" { + t.Fatalf("got %q", got) + } +} + +func TestEinoExtractFallbackAssistantFromMsgs_fromAssistantToolCalls(t *testing.T) { + m := schema.AssistantMessage("", []schema.ToolCall{{ + ID: "x", + Type: "function", + Function: schema.FunctionCall{ + Name: "exit", + Arguments: `{"final_result":"from args"}`, + }, + }}) + if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{m}); got != "from args" { + t.Fatalf("got %q", got) + } +} + +func TestEinoExtractFallbackAssistantFromMsgs_prefersToolOverEarlierAssistant(t *testing.T) { + asst := schema.AssistantMessage("", []schema.ToolCall{{ + ID: "x", + Type: "function", + Function: schema.FunctionCall{ + Name: "exit", + Arguments: `{"final_result":"from args"}`, + }, + }}) + tool := toolExitMsg("from tool", "c1") + if got := einoExtractFallbackAssistantFromMsgs([]*schema.Message{asst, tool}); got != "from tool" { + t.Fatalf("got %q", got) + } +} + +func toolExitMsg(content, callID string) *schema.Message { + m := schema.ToolMessage(content, callID) + m.ToolName = "exit" + return m +} diff --git a/internal/multiagent/eino_filesystem_tool_monitor.go b/internal/multiagent/eino_filesystem_tool_monitor.go new file mode 100644 index 00000000..8e3a8cb5 --- /dev/null +++ b/internal/multiagent/eino_filesystem_tool_monitor.go @@ -0,0 +1,157 @@ +package multiagent + +import ( + "context" + "encoding/json" + "errors" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/einomcp" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// einoADKFilesystemToolNames 与 cloudwego/eino/adk/middlewares/filesystem 默认 ToolName* 一致。 +// execute 已由 eino_execute_monitor 落库,此处不包含。 +var einoADKFilesystemToolNames = map[string]struct{}{ + "ls": {}, + "read_file": {}, + "write_file": {}, + "edit_file": {}, + "glob": {}, + "grep": {}, +} + +func isBuiltinEinoADKFilesystemToolName(name string) bool { + n := strings.ToLower(strings.TrimSpace(name)) + _, ok := einoADKFilesystemToolNames[n] + return ok +} + +func toolCallArgsFromAccumulated(msgs []adk.Message, toolCallID, expectToolName string) map[string]interface{} { + tid := strings.TrimSpace(toolCallID) + expect := strings.TrimSpace(expectToolName) + for i := len(msgs) - 1; i >= 0; i-- { + m := msgs[i] + if m == nil || m.Role != schema.Assistant || len(m.ToolCalls) == 0 { + continue + } + for j := len(m.ToolCalls) - 1; j >= 0; j-- { + tc := m.ToolCalls[j] + if tid != "" && strings.TrimSpace(tc.ID) != tid { + continue + } + fn := strings.TrimSpace(tc.Function.Name) + if expect != "" && !strings.EqualFold(fn, expect) { + continue + } + raw := strings.TrimSpace(tc.Function.Arguments) + if raw == "" { + return map[string]interface{}{} + } + var args map[string]interface{} + if err := json.Unmarshal([]byte(raw), &args); err != nil { + return map[string]interface{}{"arguments_raw": raw} + } + if args == nil { + return map[string]interface{}{} + } + return args + } + } + return map[string]interface{}{} +} + +func mustMarshalToolArguments(args map[string]interface{}) string { + if len(args) == 0 { + return "{}" + } + raw, err := json.Marshal(args) + if err != nil { + return "{}" + } + return string(raw) +} + +// beginEinoADKFilesystemToolMonitor 在 Eino ADK filesystem 工具开始调用时写入 running 状态。 +func beginEinoADKFilesystemToolMonitor( + ctx context.Context, + ag *agent.Agent, + rec einomcp.ExecutionRecorder, + binder *MCPExecutionBinder, + toolCallID, toolName string, + args map[string]interface{}, +) { + if ag == nil || rec == nil { + return + } + name := strings.TrimSpace(toolName) + if name == "" || strings.EqualFold(name, "execute") { + return + } + if !isBuiltinEinoADKFilesystemToolName(name) { + return + } + tid := strings.TrimSpace(toolCallID) + if tid == "" { + return + } + storedName := "eino_fs::" + strings.ToLower(name) + id := ag.BeginLocalToolExecution(ctx, storedName, args) + if id == "" { + return + } + rec(id, tid) + if binder != nil { + binder.Bind(tid, id) + } +} + +// recordEinoADKFilesystemToolMonitor 将 Eino ADK filesystem 中间件工具结果写入 MCP 监控(与 execute / MCP 桥芯片一致)。 +func recordEinoADKFilesystemToolMonitor( + ctx context.Context, + ag *agent.Agent, + rec einomcp.ExecutionRecorder, + binder *MCPExecutionBinder, + toolName string, + toolCallID string, + msgs []adk.Message, + resultText string, + isErr bool, +) string { + if ag == nil || rec == nil { + return "" + } + name := strings.TrimSpace(toolName) + if name == "" || strings.EqualFold(name, "execute") { + return "" + } + if !isBuiltinEinoADKFilesystemToolName(name) { + return "" + } + args := toolCallArgsFromAccumulated(msgs, toolCallID, name) + if len(args) == 0 && binder != nil { + args = binder.Arguments(toolCallID) + } + storedName := "eino_fs::" + strings.ToLower(name) + var invErr error + if isErr { + t := strings.TrimSpace(resultText) + if t == "" { + invErr = errors.New("tool error") + } else { + invErr = errors.New(t) + } + } + execID := "" + if binder != nil { + execID = binder.ExecutionID(toolCallID) + } + id := ag.FinishLocalToolExecution(ctx, execID, storedName, args, resultText, invErr) + if id != "" && execID == "" { + rec(id, toolCallID) + } + return id +} diff --git a/internal/multiagent/eino_filesystem_tool_monitor_test.go b/internal/multiagent/eino_filesystem_tool_monitor_test.go new file mode 100644 index 00000000..1beaa6ae --- /dev/null +++ b/internal/multiagent/eino_filesystem_tool_monitor_test.go @@ -0,0 +1,170 @@ +package multiagent + +import ( + "context" + "strings" + "testing" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/mcp" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +func TestEinoADKFilesystemToolMonitorBindsFinishesAndUpdatesDisplayResult(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := zap.NewNop() + server := mcp.NewServer(logger) + ag := agent.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 1) + binder := NewMCPExecutionBinder() + var recorded []string + rec := einomcp.ExecutionRecorder(func(executionID, toolCallID string) { + recorded = append(recorded, executionID+"|"+toolCallID) + }) + + beginEinoADKFilesystemToolMonitor(ctx, ag, rec, binder, "call-read", "read_file", map[string]interface{}{"path": "/tmp/secret.txt"}) + execID := binder.ExecutionID("call-read") + if execID == "" { + t.Fatal("expected begin to bind execution id") + } + exec, ok := server.GetExecution(execID) + if !ok || exec == nil || exec.Status != "running" || exec.ToolName != "eino_fs::read_file" { + t.Fatalf("begin execution = %#v ok=%v", exec, ok) + } + if len(recorded) != 1 || recorded[0] != execID+"|call-read" { + t.Fatalf("recorded begin ids = %#v", recorded) + } + if got, _ := exec.Arguments["path"].(string); got != "/tmp/secret.txt" { + t.Fatalf("begin execution args = %#v", exec.Arguments) + } + + runMessages := newEinoRunMessageAccumulator([]adk.Message{ + &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call-read", + Type: "function", + Function: schema.FunctionCall{ + Name: "read_file", + Arguments: `{"path":"/tmp/secret.txt"}`, + }, + }}, + }, + }) + emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{ + ConversationID: "conv-1", + RunMessages: runMessages, + FilesystemMonitorAgent: ag, + FilesystemMonitorRecord: rec, + MCPExecutionBinder: binder, + }) + + if !emitter.Emit(ctx, "read_file", "model-facing truncated body", "call-read", false, "lead") { + t.Fatal("expected tool_result emit") + } + exec, ok = server.GetExecution(execID) + if !ok || exec == nil { + t.Fatalf("finished execution missing: ok=%v exec=%#v", ok, exec) + } + if exec.Status != "completed" || exec.ToolName != "eino_fs::read_file" { + t.Fatalf("finished execution status/name = %#v", exec) + } + if got, _ := exec.Arguments["path"].(string); got != "/tmp/secret.txt" { + t.Fatalf("execution args = %#v", exec.Arguments) + } + if exec.Result == nil || len(exec.Result.Content) != 1 || exec.Result.Content[0].Text != "model-facing truncated body" { + t.Fatalf("execution display result = %#v", exec.Result) + } + if len(recorded) != 1 { + t.Fatalf("finish should reuse existing execution without recording a second id, got %#v", recorded) + } +} + +func TestEinoADKFilesystemToolMonitorSpillsLargeReadFileResultForProgress(t *testing.T) { + t.Parallel() + ctx := context.Background() + logger := zap.NewNop() + server := mcp.NewServer(logger) + server.ConfigureToolResultMaxBytes(400) + server.ConfigureToolResultSpillRoot(t.TempDir()) + ag := agent.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 1) + binder := NewMCPExecutionBinder() + rec := einomcp.ExecutionRecorder(func(executionID, toolCallID string) {}) + var event map[string]interface{} + + runMessages := newEinoRunMessageAccumulator([]adk.Message{ + &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call-read", + Type: "function", + Function: schema.FunctionCall{ + Name: "read_file", + Arguments: `{"path":"/tmp/large.txt"}`, + }, + }}, + }, + }) + beginEinoADKFilesystemToolMonitor(ctx, ag, rec, binder, "call-read", "read_file", map[string]interface{}{"path": "/tmp/large.txt"}) + emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{ + ConversationID: "conv-1", + RunMessages: runMessages, + FilesystemMonitorAgent: ag, + FilesystemMonitorRecord: rec, + MCPExecutionBinder: binder, + Progress: func(eventType, _ string, data interface{}) { + if eventType == "tool_result" { + event, _ = data.(map[string]interface{}) + } + }, + }) + + if !emitter.Emit(ctx, "read_file", strings.Repeat("0123456789", 100), "call-read", false, "lead") { + t.Fatal("expected tool result emit") + } + result, _ := event["result"].(string) + if !strings.Contains(result, "") || !strings.Contains(result, "Full output saved to:") { + t.Fatalf("large read_file result was not spilled in progress event: %q", result) + } + if len(result) > 400 { + t.Fatalf("progress result exceeded configured max: len=%d text=%q", len(result), result) + } +} + +func TestEinoAgenticFilesystemWrapperCapturesArgumentsAndSpillsResult(t *testing.T) { + t.Parallel() + binder := NewMCPExecutionBinder() + mw := &einoAgenticFilesystemToolMiddleware{ + TypedChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + conversationID: "conv-1", + toolMaxBytes: 400, + reductionRootDir: t.TempDir(), + binder: binder, + } + endpoint, err := mw.WrapInvokableToolCall(context.Background(), func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + return strings.Repeat("0123456789", 100), nil + }, &adk.ToolContext{Name: "read_file", CallID: "call-read"}) + if err != nil { + t.Fatalf("WrapInvokableToolCall: %v", err) + } + result, err := endpoint(context.Background(), `{"file_path":"/tmp/requirements.txt","limit":2000}`) + if err != nil { + t.Fatalf("endpoint: %v", err) + } + args := binder.Arguments("call-read") + if args["file_path"] != "/tmp/requirements.txt" { + t.Fatalf("captured args = %#v", args) + } + if !strings.Contains(result, "") || !strings.Contains(result, "Full output saved to:") { + t.Fatalf("expected persisted-output summary, got %q", result) + } + if len(result) > 400 { + t.Fatalf("summary exceeded max bytes: len=%d", len(result)) + } +} diff --git a/internal/multiagent/eino_initial_iterator_start_handler.go b/internal/multiagent/eino_initial_iterator_start_handler.go new file mode 100644 index 00000000..630b6a14 --- /dev/null +++ b/internal/multiagent/eino_initial_iterator_start_handler.go @@ -0,0 +1,54 @@ +package multiagent + +import "github.com/cloudwego/eino/adk" + +type einoAgentEventIteratorStarter func([]adk.Message) *adk.AsyncIterator[*adk.AgentEvent] + +type einoInitialIteratorStartHandlerConfig struct { + ConversationID string + OrchMode string + Progress func(eventType, message string, data interface{}) + UseTurnLoop bool + StartRunner einoAgentEventIteratorStarter + StartTurnLoop einoAgentEventIteratorStarter +} + +type einoInitialIteratorStartHandler struct { + cfg einoInitialIteratorStartHandlerConfig +} + +func newEinoInitialIteratorStartHandler(cfg einoInitialIteratorStartHandlerConfig) *einoInitialIteratorStartHandler { + return &einoInitialIteratorStartHandler{cfg: cfg} +} + +func (h *einoInitialIteratorStartHandler) StartIfNeeded(existing *adk.AsyncIterator[*adk.AgentEvent], msgs []adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + if existing != nil { + return existing + } + if h == nil { + return nil + } + if h.cfg.UseTurnLoop { + h.emitTurnLoopTakeover() + if h.cfg.StartTurnLoop == nil { + return nil + } + return h.cfg.StartTurnLoop(msgs) + } + if h.cfg.StartRunner == nil { + return nil + } + return h.cfg.StartRunner(msgs) +} + +func (h *einoInitialIteratorStartHandler) emitTurnLoopTakeover() { + if h == nil || h.cfg.Progress == nil { + return + } + h.cfg.Progress("progress", "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。", map[string]interface{}{ + "conversationId": h.cfg.ConversationID, + "source": "eino", + "orchestration": h.cfg.OrchMode, + "kind": "turn_loop_takeover", + }) +} diff --git a/internal/multiagent/eino_initial_iterator_start_handler_test.go b/internal/multiagent/eino_initial_iterator_start_handler_test.go new file mode 100644 index 00000000..3564e3e1 --- /dev/null +++ b/internal/multiagent/eino_initial_iterator_start_handler_test.go @@ -0,0 +1,111 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/adk" +) + +func TestEinoInitialIteratorStartHandlerKeepsExistingIterator(t *testing.T) { + existing, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + defer gen.Close() + + var started bool + got := newEinoInitialIteratorStartHandler(einoInitialIteratorStartHandlerConfig{ + UseTurnLoop: true, + StartTurnLoop: func([]adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + started = true + iter, iterGen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + iterGen.Close() + return iter + }, + Progress: func(string, string, interface{}) { + t.Fatal("progress should not be emitted when an iterator already exists") + }, + }).StartIfNeeded(existing, nil) + + if got != existing { + t.Fatal("existing iterator should be preserved") + } + if started { + t.Fatal("start function should not be called when an iterator already exists") + } +} + +func TestEinoInitialIteratorStartHandlerStartsRunner(t *testing.T) { + wantIter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + defer gen.Close() + + var runnerStarted bool + got := newEinoInitialIteratorStartHandler(einoInitialIteratorStartHandlerConfig{ + StartRunner: func(msgs []adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + runnerStarted = true + if msgs == nil { + t.Fatal("msgs should be forwarded") + } + return wantIter + }, + StartTurnLoop: func([]adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + t.Fatal("turn loop should not start when UseTurnLoop is false") + return nil + }, + Progress: func(string, string, interface{}) { + t.Fatal("runner start should not emit TurnLoop takeover progress") + }, + }).StartIfNeeded(nil, []adk.Message{}) + + if !runnerStarted { + t.Fatal("runner start was not called") + } + if got != wantIter { + t.Fatal("runner iterator should be returned") + } +} + +func TestEinoInitialIteratorStartHandlerStartsTurnLoopWithTakeoverProgress(t *testing.T) { + wantIter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + defer gen.Close() + + var turnLoopStarted bool + var gotType, gotMessage string + var gotData map[string]interface{} + got := newEinoInitialIteratorStartHandler(einoInitialIteratorStartHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + UseTurnLoop: true, + StartRunner: func([]adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + t.Fatal("runner should not start when UseTurnLoop is true") + return nil + }, + StartTurnLoop: func(msgs []adk.Message) *adk.AsyncIterator[*adk.AgentEvent] { + turnLoopStarted = true + if msgs == nil { + t.Fatal("msgs should be forwarded") + } + return wantIter + }, + Progress: func(eventType, message string, data interface{}) { + gotType = eventType + gotMessage = message + if m, ok := data.(map[string]interface{}); ok { + gotData = m + } + }, + }).StartIfNeeded(nil, []adk.Message{}) + + if !turnLoopStarted { + t.Fatal("turn loop start was not called") + } + if got != wantIter { + t.Fatal("turn loop iterator should be returned") + } + if gotType != "progress" { + t.Fatalf("progress type = %q, want progress", gotType) + } + if gotMessage != "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。" { + t.Fatalf("progress message = %q", gotMessage) + } + if gotData["conversationId"] != "conv-1" || gotData["source"] != "eino" || gotData["orchestration"] != "deep" { + t.Fatalf("progress data = %#v", gotData) + } +} diff --git a/internal/multiagent/eino_input_telemetry.go b/internal/multiagent/eino_input_telemetry.go new file mode 100644 index 00000000..dbf3c576 --- /dev/null +++ b/internal/multiagent/eino_input_telemetry.go @@ -0,0 +1,133 @@ +package multiagent + +import ( + "context" + "strings" + + "cyberstrike-ai/internal/agent" + + "github.com/bytedance/sonic" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +type einoModelInputTelemetryMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + modelName string + conversationID string + phase string +} + +func newEinoModelInputTelemetryMiddleware( + logger *zap.Logger, + modelName string, + conversationID string, + phase string, +) adk.ChatModelAgentMiddleware { + if logger == nil { + return nil + } + return &einoModelInputTelemetryMiddleware{ + logger: logger, + modelName: strings.TrimSpace(modelName), + conversationID: strings.TrimSpace(conversationID), + phase: strings.TrimSpace(phase), + } +} + +func (m *einoModelInputTelemetryMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.ChatModelAgentState, + mc *adk.ModelContext, +) (context.Context, *adk.ChatModelAgentState, error) { + if m == nil || m.logger == nil || state == nil { + return ctx, state, nil + } + tokens := estimateTokensForMessagesAndTools(ctx, m.modelName, state.Messages, mcTools(mc)) + m.logger.Info("eino model input estimated", + zap.String("phase", m.phase), + zap.String("conversation_id", m.conversationID), + zap.Int("messages", len(state.Messages)), + zap.Int("tools", len(mcTools(mc))), + zap.Int("input_tokens_estimated", tokens), + ) + return ctx, state, nil +} + +func mcTools(mc *adk.ModelContext) []*schema.ToolInfo { + if mc == nil || len(mc.Tools) == 0 { + return nil + } + return mc.Tools +} + +func estimateTokensForMessagesAndTools( + _ context.Context, + modelName string, + messages []adk.Message, + tools []*schema.ToolInfo, +) int { + var sb strings.Builder + for _, msg := range messages { + if msg == nil { + continue + } + sb.WriteString(string(msg.Role)) + sb.WriteByte('\n') + sb.WriteString(msg.Content) + sb.WriteByte('\n') + if msg.ReasoningContent != "" { + sb.WriteString(msg.ReasoningContent) + sb.WriteByte('\n') + } + if len(msg.ToolCalls) > 0 { + if b, err := sonic.Marshal(msg.ToolCalls); err == nil { + sb.Write(b) + sb.WriteByte('\n') + } + } + } + for _, tl := range tools { + if tl == nil { + continue + } + cp := *tl + cp.Extra = nil + if text, err := sonic.MarshalString(cp); err == nil { + sb.WriteString(text) + sb.WriteByte('\n') + } + } + text := sb.String() + if text == "" { + return 0 + } + tc := agent.NewTikTokenCounter() + if n, err := tc.Count(modelName, text); err == nil { + return n + } + return (len(text) + 3) / 4 +} + +func logPlanExecuteModelInputEstimate( + logger *zap.Logger, + modelName string, + conversationID string, + phase string, + msgs []adk.Message, +) { + if logger == nil { + return + } + tokens := estimateTokensForMessagesAndTools(context.Background(), modelName, msgs, nil) + logger.Info("eino model input estimated", + zap.String("phase", phase), + zap.String("conversation_id", strings.TrimSpace(conversationID)), + zap.Int("messages", len(msgs)), + zap.Int("tools", 0), + zap.Int("input_tokens_estimated", tokens), + ) +} + diff --git a/internal/multiagent/eino_main_assistant_complete_handler.go b/internal/multiagent/eino_main_assistant_complete_handler.go new file mode 100644 index 00000000..427ee854 --- /dev/null +++ b/internal/multiagent/eino_main_assistant_complete_handler.go @@ -0,0 +1,49 @@ +package multiagent + +import "strings" + +type einoMainAssistantCompleteHandler struct { + agentName string + emitter *einoMainResponseStreamEmitter + stdoutSuppressor *einoExecuteStdoutSuppressor + assistantOutput *einoAssistantOutputAccumulator +} + +type einoMainAssistantCompleteHandlerConfig struct { + AgentName string + Emitter *einoMainResponseStreamEmitter + StdoutSuppressor *einoExecuteStdoutSuppressor + AssistantOutput *einoAssistantOutputAccumulator +} + +func newEinoMainAssistantCompleteHandler(cfg einoMainAssistantCompleteHandlerConfig) *einoMainAssistantCompleteHandler { + return &einoMainAssistantCompleteHandler{ + agentName: cfg.AgentName, + emitter: cfg.Emitter, + stdoutSuppressor: cfg.StdoutSuppressor, + assistantOutput: cfg.AssistantOutput, + } +} + +func (h *einoMainAssistantCompleteHandler) EmitComplete(content string) bool { + if h == nil { + return false + } + body := strings.TrimSpace(content) + if body == "" { + return false + } + if h.stdoutSuppressor != nil { + if dup := h.stdoutSuppressor.Consume(); dup != "" && body == dup { + if h.assistantOutput != nil { + h.assistantOutput.RecordMainAssistant(h.agentName, body) + } + return false + } + } + emitted := h.emitter.EmitDelta(body, body) + if h.assistantOutput != nil { + h.assistantOutput.RecordMainAssistant(h.agentName, body) + } + return emitted +} diff --git a/internal/multiagent/eino_main_assistant_complete_handler_test.go b/internal/multiagent/eino_main_assistant_complete_handler_test.go new file mode 100644 index 00000000..577604f5 --- /dev/null +++ b/internal/multiagent/eino_main_assistant_complete_handler_test.go @@ -0,0 +1,76 @@ +package multiagent + +import "testing" + +func TestEinoMainAssistantCompleteHandlerEmitsAndRecords(t *testing.T) { + var eventTypes []string + var messages []string + progress := func(eventType, message string, _ interface{}) { + eventTypes = append(eventTypes, eventType) + messages = append(messages, message) + } + out := newEinoAssistantOutputAccumulator("deep") + handler := newEinoMainAssistantCompleteHandler(einoMainAssistantCompleteHandlerConfig{ + AgentName: "lead", + Emitter: newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 2, progress, nil), + AssistantOutput: out, + }) + + if !handler.EmitComplete(" hello ") { + t.Fatal("complete assistant should emit") + } + if len(eventTypes) != 2 || eventTypes[0] != "response_start" || eventTypes[1] != "response_delta" { + t.Fatalf("events = %#v", eventTypes) + } + if messages[1] != "hello" { + t.Fatalf("delta message = %q", messages[1]) + } + if out.LastAssistant() != "hello" { + t.Fatalf("last assistant = %q", out.LastAssistant()) + } +} + +func TestEinoMainAssistantCompleteHandlerSuppressesDuplicateExecuteStdout(t *testing.T) { + var eventTypes []string + progress := func(eventType, _ string, _ interface{}) { + eventTypes = append(eventTypes, eventType) + } + stdoutDup := newEinoExecuteStdoutSuppressor() + stdoutDup.Record("execute", "hello", false) + out := newEinoAssistantOutputAccumulator("deep") + handler := newEinoMainAssistantCompleteHandler(einoMainAssistantCompleteHandlerConfig{ + AgentName: "lead", + Emitter: newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 1, progress, nil), + StdoutSuppressor: stdoutDup, + AssistantOutput: out, + }) + + if handler.EmitComplete("hello") { + t.Fatal("duplicate execute stdout should not emit") + } + if len(eventTypes) != 0 { + t.Fatalf("events = %#v, want none", eventTypes) + } + if out.LastAssistant() != "hello" { + t.Fatalf("last assistant = %q", out.LastAssistant()) + } + if stdoutDup.Peek() != "" { + t.Fatal("duplicate target should be consumed") + } +} + +func TestEinoMainAssistantCompleteHandlerRecordsWithoutProgress(t *testing.T) { + out := newEinoAssistantOutputAccumulator("plan_execute") + handler := newEinoMainAssistantCompleteHandler(einoMainAssistantCompleteHandlerConfig{ + AgentName: "executor", + Emitter: newEinoMainResponseStreamEmitter("conv-1", "plan_execute", "executor", "stream-1", 1, nil, nil), + AssistantOutput: out, + }) + + if handler.EmitComplete(`{"response":"done"}`) { + t.Fatal("nil progress should not emit") + } + if out.LastPlanExecuteExecutor() != "done" { + t.Fatalf("executor output = %q", out.LastPlanExecuteExecutor()) + } +} diff --git a/internal/multiagent/eino_main_assistant_stream_handler.go b/internal/multiagent/eino_main_assistant_stream_handler.go new file mode 100644 index 00000000..8e6d9207 --- /dev/null +++ b/internal/multiagent/eino_main_assistant_stream_handler.go @@ -0,0 +1,77 @@ +package multiagent + +import "strings" + +type einoMainAssistantStreamHandler struct { + agentName string + emitter *einoMainResponseStreamEmitter + stdoutSuppressor *einoExecuteStdoutSuppressor + assistantOutput *einoAssistantOutputAccumulator + runMessages *einoRunMessageAccumulator + + buf string + dupTarget string +} + +type einoMainAssistantStreamHandlerConfig struct { + AgentName string + Emitter *einoMainResponseStreamEmitter + StdoutSuppressor *einoExecuteStdoutSuppressor + AssistantOutput *einoAssistantOutputAccumulator + RunMessages *einoRunMessageAccumulator +} + +func newEinoMainAssistantStreamHandler(cfg einoMainAssistantStreamHandlerConfig) *einoMainAssistantStreamHandler { + return &einoMainAssistantStreamHandler{ + agentName: cfg.AgentName, + emitter: cfg.Emitter, + stdoutSuppressor: cfg.StdoutSuppressor, + assistantOutput: cfg.AssistantOutput, + runMessages: cfg.RunMessages, + } +} + +func (h *einoMainAssistantStreamHandler) EmitDelta(content string) bool { + if h == nil || content == "" { + return false + } + var delta string + h.buf, delta = normalizeStreamingDelta(h.buf, content) + if delta == "" { + return false + } + if h.dupTarget == "" && h.stdoutSuppressor != nil { + h.dupTarget = h.stdoutSuppressor.Peek() + } + if h.dupTarget != "" { + return false + } + return h.emitter.EmitDelta(delta, h.buf) +} + +func (h *einoMainAssistantStreamHandler) Finish() string { + if h == nil { + return "" + } + body := strings.TrimSpace(h.buf) + if body == "" { + return "" + } + if h.dupTarget != "" { + if h.stdoutSuppressor != nil { + h.stdoutSuppressor.Clear() + } + if body != h.dupTarget { + h.emitter.EmitTailFromFull(h.buf) + } + } else { + h.emitter.EmitTailFromFull(h.buf) + } + if h.assistantOutput != nil { + h.assistantOutput.RecordMainAssistant(h.agentName, body) + } + if h.runMessages != nil { + h.runMessages.AppendAssistantText(body) + } + return body +} diff --git a/internal/multiagent/eino_main_assistant_stream_handler_test.go b/internal/multiagent/eino_main_assistant_stream_handler_test.go new file mode 100644 index 00000000..5feb7adf --- /dev/null +++ b/internal/multiagent/eino_main_assistant_stream_handler_test.go @@ -0,0 +1,103 @@ +package multiagent + +import "testing" + +func TestEinoMainAssistantStreamHandlerEmitsAndRecords(t *testing.T) { + var eventTypes []string + var messages []string + progress := func(eventType, message string, _ interface{}) { + eventTypes = append(eventTypes, eventType) + messages = append(messages, message) + } + out := newEinoAssistantOutputAccumulator("deep") + runMsgs := newEinoRunMessageAccumulator(nil) + emitter := newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 2, progress, nil) + handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{ + AgentName: "lead", + Emitter: emitter, + AssistantOutput: out, + RunMessages: runMsgs, + }) + + if !handler.EmitDelta("he") { + t.Fatal("first delta should emit") + } + if !handler.EmitDelta("hello") { + t.Fatal("cumulative chunk should emit tail") + } + if got := handler.Finish(); got != "hello" { + t.Fatalf("finish = %q, want hello", got) + } + + if len(eventTypes) != 3 || eventTypes[0] != "response_start" || eventTypes[1] != "response_delta" || eventTypes[2] != "response_delta" { + t.Fatalf("events = %#v", eventTypes) + } + if messages[1] != "he" || messages[2] != "llo" { + t.Fatalf("delta messages = %#v", messages) + } + if out.LastAssistant() != "hello" { + t.Fatalf("last assistant = %q", out.LastAssistant()) + } + if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != "hello" { + t.Fatalf("run messages = %#v", msgs) + } +} + +func TestEinoMainAssistantStreamHandlerSuppressesDuplicateExecuteStdout(t *testing.T) { + var eventTypes []string + progress := func(eventType, _ string, _ interface{}) { + eventTypes = append(eventTypes, eventType) + } + stdoutDup := newEinoExecuteStdoutSuppressor() + stdoutDup.Record("execute", "hello", false) + out := newEinoAssistantOutputAccumulator("deep") + runMsgs := newEinoRunMessageAccumulator(nil) + handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{ + AgentName: "lead", + Emitter: newEinoMainResponseStreamEmitter("conv-1", "deep", "lead", "stream-1", 1, progress, nil), + StdoutSuppressor: stdoutDup, + AssistantOutput: out, + RunMessages: runMsgs, + }) + + if handler.EmitDelta("hello") { + t.Fatal("duplicate execute stdout should not emit delta") + } + if got := handler.Finish(); got != "hello" { + t.Fatalf("finish = %q, want hello", got) + } + if len(eventTypes) != 0 { + t.Fatalf("events = %#v, want none", eventTypes) + } + if stdoutDup.Peek() != "" { + t.Fatal("duplicate target should be cleared on finish") + } + if out.LastAssistant() != "hello" { + t.Fatalf("last assistant = %q", out.LastAssistant()) + } + if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != "hello" { + t.Fatalf("run messages = %#v", msgs) + } +} + +func TestEinoMainAssistantStreamHandlerRecordsWithoutProgress(t *testing.T) { + out := newEinoAssistantOutputAccumulator("plan_execute") + runMsgs := newEinoRunMessageAccumulator(nil) + handler := newEinoMainAssistantStreamHandler(einoMainAssistantStreamHandlerConfig{ + AgentName: "executor", + Emitter: newEinoMainResponseStreamEmitter("conv-1", "plan_execute", "executor", "stream-1", 1, nil, nil), + AssistantOutput: out, + RunMessages: runMsgs, + }) + + handler.EmitDelta(`{"response":"done"}`) + if got := handler.Finish(); got != `{"response":"done"}` { + t.Fatalf("finish = %q", got) + } + if out.LastPlanExecuteExecutor() != "done" { + t.Fatalf("executor output = %q", out.LastPlanExecuteExecutor()) + } + if msgs := runMsgs.Messages(); len(msgs) != 1 || msgs[0].Content != `{"response":"done"}` { + t.Fatalf("run messages = %#v", msgs) + } +} diff --git a/internal/multiagent/eino_main_response_stream_emitter.go b/internal/multiagent/eino_main_response_stream_emitter.go new file mode 100644 index 00000000..19dcc69d --- /dev/null +++ b/internal/multiagent/eino_main_response_stream_emitter.go @@ -0,0 +1,85 @@ +package multiagent + +import "cyberstrike-ai/internal/openai" + +type einoMainResponseStreamEmitter struct { + progress func(eventType, message string, data interface{}) + snapshotMCPIDs func() []string + conversationID string + orchMode string + agentName string + streamID string + iteration int + headerSent bool + wireAccum string +} + +func newEinoMainResponseStreamEmitter( + conversationID, orchMode, agentName, streamID string, + iteration int, + progress func(eventType, message string, data interface{}), + snapshotMCPIDs func() []string, +) *einoMainResponseStreamEmitter { + if snapshotMCPIDs == nil { + snapshotMCPIDs = func() []string { return nil } + } + return &einoMainResponseStreamEmitter{ + progress: progress, + snapshotMCPIDs: snapshotMCPIDs, + conversationID: conversationID, + orchMode: orchMode, + agentName: agentName, + streamID: streamID, + iteration: iteration, + } +} + +func (e *einoMainResponseStreamEmitter) EmitDelta(delta, accumulated string) bool { + if e == nil || e.progress == nil || delta == "" { + return false + } + e.emitStart() + e.progress("response_delta", delta, openai.WithSSEAccumulated(e.responseData(), accumulated)) + e.wireAccum, _ = normalizeStreamingDelta(e.wireAccum, delta) + return true +} + +func (e *einoMainResponseStreamEmitter) EmitTailFromFull(full string) bool { + if e == nil || full == "" { + return false + } + _, tail := normalizeStreamingDelta(e.wireAccum, full) + if tail == "" { + return false + } + return e.EmitDelta(tail, full) +} + +func (e *einoMainResponseStreamEmitter) emitStart() { + if e.headerSent || e.progress == nil { + return + } + e.progress("response_start", "", map[string]interface{}{ + "conversationId": e.conversationID, + "mcpExecutionIds": e.snapshotMCPIDs(), + "messageGeneratedBy": "eino:" + e.agentName, + "einoRole": "orchestrator", + "einoAgent": e.agentName, + "orchestration": e.orchMode, + "iteration": e.iteration, + "streamId": e.streamID, + }) + e.headerSent = true +} + +func (e *einoMainResponseStreamEmitter) responseData() map[string]interface{} { + return map[string]interface{}{ + "conversationId": e.conversationID, + "mcpExecutionIds": e.snapshotMCPIDs(), + "einoRole": "orchestrator", + "einoAgent": e.agentName, + "orchestration": e.orchMode, + "iteration": e.iteration, + "streamId": e.streamID, + } +} diff --git a/internal/multiagent/eino_main_response_stream_emitter_test.go b/internal/multiagent/eino_main_response_stream_emitter_test.go new file mode 100644 index 00000000..36188bd0 --- /dev/null +++ b/internal/multiagent/eino_main_response_stream_emitter_test.go @@ -0,0 +1,65 @@ +package multiagent + +import ( + "testing" + + "cyberstrike-ai/internal/openai" +) + +func TestEinoMainResponseStreamEmitterEmitsStartOnceAndTail(t *testing.T) { + type progressEvent struct { + eventType string + message string + data map[string]interface{} + } + var events []progressEvent + progress := func(eventType, message string, data interface{}) { + m, _ := data.(map[string]interface{}) + events = append(events, progressEvent{eventType: eventType, message: message, data: m}) + } + + emitter := newEinoMainResponseStreamEmitter( + "conv-1", "supervisor", "lead", "stream-1", 3, progress, func() []string { return []string{"mcp-1"} }, + ) + if !emitter.EmitDelta("he", "he") { + t.Fatal("first delta should be emitted") + } + if !emitter.EmitTailFromFull("hello") { + t.Fatal("tail should be emitted") + } + if emitter.EmitTailFromFull("hello") { + t.Fatal("duplicate tail should not be emitted") + } + + if len(events) != 3 { + t.Fatalf("events = %#v, want start + 2 deltas", events) + } + if events[0].eventType != "response_start" { + t.Fatalf("event[0] = %s, want response_start", events[0].eventType) + } + if events[1].eventType != "response_delta" || events[1].message != "he" { + t.Fatalf("event[1] = %#v, want first delta", events[1]) + } + if events[2].eventType != "response_delta" || events[2].message != "llo" { + t.Fatalf("event[2] = %#v, want tail delta", events[2]) + } + if got := events[2].data[openai.SSEAccumulatedKey]; got != "hello" { + t.Fatalf("accumulated = %#v, want hello", got) + } + if got := events[0].data["messageGeneratedBy"]; got != "eino:lead" { + t.Fatalf("messageGeneratedBy = %#v", got) + } + if got := events[0].data["iteration"]; got != 3 { + t.Fatalf("iteration = %#v", got) + } +} + +func TestEinoMainResponseStreamEmitterNoProgress(t *testing.T) { + emitter := newEinoMainResponseStreamEmitter("conv", "deep", "agent", "stream", 1, nil, nil) + if emitter.EmitDelta("hello", "hello") { + t.Fatal("nil progress should not emit") + } + if emitter.EmitTailFromFull("hello") { + t.Fatal("nil progress should not emit tail") + } +} diff --git a/internal/multiagent/eino_materialized_message_event_handler.go b/internal/multiagent/eino_materialized_message_event_handler.go new file mode 100644 index 00000000..810de05d --- /dev/null +++ b/internal/multiagent/eino_materialized_message_event_handler.go @@ -0,0 +1,115 @@ +package multiagent + +import ( + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +type einoMaterializedMessageEventHandlerConfig struct { + ConversationID string + OrchMode string + Progress func(eventType, message string, data interface{}) + SnapshotMCPIDs func() []string + StreamsMainAssistant func(agent string) bool + EinoRoleTag func(agent string) string + RunProgress *einoRunProgressTracker + StdoutSuppressor *einoExecuteStdoutSuppressor + AssistantOutput *einoAssistantOutputAccumulator + RunMessages *einoRunMessageAccumulator + Usage *einoRunUsageAccumulator + ToolResultHandler *einoToolResultEventHandler + MarkPending func(toolCallPendingInfo) + NextMainStreamID func() string +} + +type einoMaterializedMessageEventHandler struct { + conversationID string + orchMode string + progress func(eventType, message string, data interface{}) + snapshotMCPIDs func() []string + streamsMainAssistant func(agent string) bool + einoRoleTag func(agent string) string + runProgress *einoRunProgressTracker + stdoutSuppressor *einoExecuteStdoutSuppressor + assistantOutput *einoAssistantOutputAccumulator + runMessages *einoRunMessageAccumulator + usage *einoRunUsageAccumulator + toolResultHandler *einoToolResultEventHandler + markPending func(toolCallPendingInfo) + nextMainStreamID func() string +} + +func newEinoMaterializedMessageEventHandler(cfg einoMaterializedMessageEventHandlerConfig) *einoMaterializedMessageEventHandler { + if cfg.SnapshotMCPIDs == nil { + cfg.SnapshotMCPIDs = func() []string { return nil } + } + if cfg.StreamsMainAssistant == nil { + cfg.StreamsMainAssistant = func(string) bool { return true } + } + if cfg.EinoRoleTag == nil { + cfg.EinoRoleTag = func(string) string { return "" } + } + if cfg.NextMainStreamID == nil { + cfg.NextMainStreamID = func() string { return "eino-main" } + } + return &einoMaterializedMessageEventHandler{ + conversationID: cfg.ConversationID, + orchMode: cfg.OrchMode, + progress: cfg.Progress, + snapshotMCPIDs: cfg.SnapshotMCPIDs, + streamsMainAssistant: cfg.StreamsMainAssistant, + einoRoleTag: cfg.EinoRoleTag, + runProgress: cfg.RunProgress, + stdoutSuppressor: cfg.StdoutSuppressor, + assistantOutput: cfg.AssistantOutput, + runMessages: cfg.RunMessages, + usage: cfg.Usage, + toolResultHandler: cfg.ToolResultHandler, + markPending: cfg.MarkPending, + nextMainStreamID: cfg.NextMainStreamID, + } +} + +func (h *einoMaterializedMessageEventHandler) Handle(mv *adk.MessageVariant, msg adk.Message, agentName string) bool { + if h == nil || mv == nil || msg == nil { + return false + } + if h.runMessages != nil { + h.runMessages.Append(msg) + } + if msg.Role == schema.Assistant && h.usage != nil { + h.usage.AddMessage(msg) + } + if h.runProgress != nil { + h.runProgress.EmitToolCalls(mergeMessageToolCalls(msg), agentName, h.markPending) + } + if mv.Role == schema.Assistant { + newEinoReasoningStreamEmitter(h.conversationID, h.orchMode, agentName, h.einoRoleTag(agentName), h.progress, nil).EmitComplete(msg.ReasoningContent) + body := strings.TrimSpace(msg.Content) + if body != "" { + if h.streamsMainAssistant(agentName) { + newEinoMainAssistantCompleteHandler(einoMainAssistantCompleteHandlerConfig{ + AgentName: agentName, + Emitter: newEinoMainResponseStreamEmitter(h.conversationID, h.orchMode, agentName, h.nextMainStreamID(), h.mainIteration(agentName), h.progress, h.snapshotMCPIDs), + StdoutSuppressor: h.stdoutSuppressor, + AssistantOutput: h.assistantOutput, + }).EmitComplete(body) + } else { + newEinoSubAgentReplyEmitter(h.conversationID, agentName, h.progress, nil).EmitComplete(body) + } + } + } + if h.toolResultHandler != nil { + h.toolResultHandler.HandleMaterialized(mv, msg, agentName) + } + return true +} + +func (h *einoMaterializedMessageEventHandler) mainIteration(agentName string) int { + if h == nil || h.runProgress == nil { + return 0 + } + return h.runProgress.MainIteration(agentName) +} diff --git a/internal/multiagent/eino_materialized_message_event_handler_test.go b/internal/multiagent/eino_materialized_message_event_handler_test.go new file mode 100644 index 00000000..a67478c6 --- /dev/null +++ b/internal/multiagent/eino_materialized_message_event_handler_test.go @@ -0,0 +1,151 @@ +package multiagent + +import ( + "testing" + + "cyberstrike-ai/internal/einomcp" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestEinoMaterializedMessageEventHandlerHandlesMainAssistant(t *testing.T) { + var events []string + runMessages := newEinoRunMessageAccumulator(nil) + assistantOutput := newEinoAssistantOutputAccumulator("deep") + usage := newEinoRunUsageAccumulator() + runProgress := newEinoRunProgressTracker( + "deep", "lead", "conv-1", + func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + func(agent string) bool { return agent == "lead" }, + nil, + ) + handler := newEinoMaterializedMessageEventHandler(einoMaterializedMessageEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + Progress: func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + RunMessages: runMessages, + Usage: usage, + AssistantOutput: assistantOutput, + RunProgress: runProgress, + StreamsMainAssistant: func(agent string) bool { return agent == "lead" }, + EinoRoleTag: func(string) string { return "orchestrator" }, + NextMainStreamID: func() string { return "main-complete-1" }, + }) + msg := schema.AssistantMessage(" done ", nil) + msg.ReasoningContent = "thought" + msg.ResponseMeta = &schema.ResponseMeta{Usage: &schema.TokenUsage{ + PromptTokens: 11, + CompletionTokens: 7, + TotalTokens: 18, + }} + mv := &adk.MessageVariant{Role: schema.Assistant} + + if !handler.Handle(mv, msg, "lead") { + t.Fatal("main assistant message was not handled") + } + if assistantOutput.LastAssistant() != "done" { + t.Fatalf("last assistant = %q", assistantOutput.LastAssistant()) + } + if msgs := runMessages.Messages(); len(msgs) != 1 || msgs[0].Content != " done " { + t.Fatalf("run messages = %#v", msgs) + } + if got := usage.Summary(); got.ModelCalls != 1 || got.TotalTokens != 18 { + t.Fatalf("usage = %#v, want one assistant model call", got) + } + if !containsString(events, "reasoning_chain") || !containsString(events, "response_start") || !containsString(events, "response_delta") { + t.Fatalf("events = %#v, want reasoning and response events", events) + } +} + +func TestEinoMaterializedMessageEventHandlerHandlesSubAssistant(t *testing.T) { + var events []string + runMessages := newEinoRunMessageAccumulator(nil) + assistantOutput := newEinoAssistantOutputAccumulator("deep") + handler := newEinoMaterializedMessageEventHandler(einoMaterializedMessageEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + Progress: func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + RunMessages: runMessages, + AssistantOutput: assistantOutput, + StreamsMainAssistant: func(agent string) bool { return agent == "lead" }, + EinoRoleTag: func(string) string { return "sub" }, + }) + + if !handler.Handle(&adk.MessageVariant{Role: schema.Assistant}, schema.AssistantMessage("sub done", nil), "worker") { + t.Fatal("sub assistant message was not handled") + } + if assistantOutput.LastAssistant() != "" { + t.Fatalf("sub assistant should not update main output, got %q", assistantOutput.LastAssistant()) + } + if len(runMessages.Messages()) != 1 { + t.Fatalf("run messages = %#v, want appended original message", runMessages.Messages()) + } + if !containsString(events, "eino_agent_reply") { + t.Fatalf("events = %#v, want sub reply event", events) + } +} + +func TestEinoMaterializedMessageEventHandlerHandlesToolCallsAndToolResult(t *testing.T) { + var events []string + var marked []toolCallPendingInfo + runMessages := newEinoRunMessageAccumulator(nil) + runProgress := newEinoRunProgressTracker( + "deep", "lead", "conv-1", + func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + func(agent string) bool { return agent == "lead" }, + nil, + ) + toolResultEmitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{ + ConversationID: "conv-1", + Progress: func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + }) + toolResultHandler := newEinoToolResultEventHandler(einoToolResultEventHandlerConfig{Emitter: toolResultEmitter}) + handler := newEinoMaterializedMessageEventHandler(einoMaterializedMessageEventHandlerConfig{ + ConversationID: "conv-1", + OrchMode: "deep", + Progress: func(eventType, _ string, _ interface{}) { events = append(events, eventType) }, + RunMessages: runMessages, + RunProgress: runProgress, + ToolResultHandler: toolResultHandler, + MarkPending: func(info toolCallPendingInfo) { + marked = append(marked, info) + }, + }) + + toolCallMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call-1", + Type: "function", + Function: schema.FunctionCall{ + Name: "execute", + Arguments: `{"command":`, + }, + }}, + } + if !handler.Handle(&adk.MessageVariant{Role: schema.Assistant}, toolCallMsg, "lead") { + t.Fatal("tool call message was not handled") + } + toolMsg := schema.ToolMessage(einomcp.ToolErrorPrefix+"bad command", "call-1", schema.WithToolName("execute")) + if !handler.Handle(&adk.MessageVariant{Role: schema.Tool}, toolMsg, "lead") { + t.Fatal("tool message was not handled") + } + + if !containsString(events, "tool_call") || !containsString(events, "tool_result") || containsString(events, "model_output_rejected") { + t.Fatalf("events = %#v, want real tool_call and tool_result without model-output recovery", events) + } + if len(marked) != 1 || marked[0].ToolCallID != "call-1" || marked[0].ToolName != "execute" { + t.Fatalf("marked pending = %#v", marked) + } + if len(runMessages.Messages()) != 2 { + t.Fatalf("run messages = %#v, want assistant and tool messages", runMessages.Messages()) + } +} + +func TestEinoMaterializedMessageEventHandlerIgnoresNil(t *testing.T) { + handler := newEinoMaterializedMessageEventHandler(einoMaterializedMessageEventHandlerConfig{}) + if handler.Handle(nil, nil, "lead") { + t.Fatal("nil message should be ignored") + } +} diff --git a/internal/multiagent/eino_message_stream_receiver.go b/internal/multiagent/eino_message_stream_receiver.go new file mode 100644 index 00000000..1cddff2d --- /dev/null +++ b/internal/multiagent/eino_message_stream_receiver.go @@ -0,0 +1,61 @@ +package multiagent + +import ( + "context" + "errors" + "io" + + "github.com/cloudwego/eino/schema" +) + +// recvEinoSchemaMessageStreamWithContext consumes an Eino schema.Message stream +// and stops promptly when ctx is canceled. EOF and nil chunks are treated as a +// normal stream boundary. +func recvEinoSchemaMessageStreamWithContext( + ctx context.Context, + stream *schema.StreamReader[*schema.Message], + buffer int, + onChunk func(*schema.Message), +) error { + if stream == nil { + return nil + } + if buffer <= 0 { + buffer = 1 + } + type streamMsg struct { + chunk *schema.Message + err error + } + recvCh := make(chan streamMsg, buffer) + go func() { + defer close(recvCh) + for { + ch, rerr := stream.Recv() + recvCh <- streamMsg{chunk: ch, err: rerr} + if rerr != nil { + return + } + } + }() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case sm, ok := <-recvCh: + if !ok { + return nil + } + if errors.Is(sm.err, io.EOF) { + return nil + } + if sm.err != nil { + return sm.err + } + if sm.chunk == nil || onChunk == nil { + continue + } + onChunk(sm.chunk) + } + } +} diff --git a/internal/multiagent/eino_middleware.go b/internal/multiagent/eino_middleware.go new file mode 100644 index 00000000..1f2c057a --- /dev/null +++ b/internal/multiagent/eino_middleware.go @@ -0,0 +1,435 @@ +package multiagent + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp/builtin" + + localbk "github.com/cloudwego/eino-ext/adk/backend/local" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/middlewares/dynamictool/toolsearch" + "github.com/cloudwego/eino/adk/middlewares/patchtoolcalls" + "github.com/cloudwego/eino/adk/middlewares/plantask" + "github.com/cloudwego/eino/adk/middlewares/reduction" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// einoMWPlacement controls which optional middleware runs on orchestrator vs sub-agents. +type einoMWPlacement int + +const ( + einoMWMain einoMWPlacement = iota // Deep / Supervisor main chat agent + einoMWSub // Specialist ChatModelAgent +) + +func sanitizeEinoPathSegment(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return "default" + } + s = strings.ReplaceAll(s, string(filepath.Separator), "-") + s = strings.ReplaceAll(s, "/", "-") + s = strings.ReplaceAll(s, "\\", "-") + s = strings.ReplaceAll(s, "..", "__") + if len(s) > 180 { + s = s[:180] + } + return s +} + +func splitToolsForToolSearch(all []tool.BaseTool, alwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) { + if alwaysVisible <= 0 || len(all) <= alwaysVisible+1 { + return all, nil, false + } + return append([]tool.BaseTool(nil), all[:alwaysVisible]...), append([]tool.BaseTool(nil), all[alwaysVisible:]...), true +} + +func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbackAlwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) { + nameSet := expandAlwaysVisibleNameSet(names) + if len(nameSet) == 0 { + return splitToolsForToolSearch(all, fallbackAlwaysVisible) + } + static = make([]tool.BaseTool, 0, len(all)) + dynamic = make([]tool.BaseTool, 0, len(all)) + for _, t := range all { + if t == nil { + continue + } + info, err := t.Info(context.Background()) + name := "" + if err == nil && info != nil { + name = info.Name + } + if toolMatchesAlwaysVisible(name, nameSet) { + static = append(static, t) + continue + } + dynamic = append(dynamic, t) + } + if len(static) == 0 || len(dynamic) == 0 { + // fallback: preserve previous behavior when whitelist misses all or includes all. + return splitToolsForToolSearch(all, fallbackAlwaysVisible) + } + return static, dynamic, true +} + +func mergeAlwaysVisibleToolNames(configured []string) []string { + merged := make([]string, 0, len(configured)+32) + seen := make(map[string]struct{}, len(configured)+32) + add := func(name string) { + n := strings.TrimSpace(strings.ToLower(name)) + if n == "" { + return + } + if _, ok := seen[n]; ok { + return + } + seen[n] = struct{}{} + merged = append(merged, n) + } + for _, n := range configured { + add(n) + } + // Always include hardcoded backend builtin MCP tools from constants. + for _, n := range builtin.GetAllBuiltinTools() { + add(n) + } + return merged +} + +func reductionCacheRootDir(configuredBase, projectID, conversationID string) string { + base := strings.TrimSpace(configuredBase) + if base == "" { + base = filepath.Join("tmp", "reduction") + } + if pid := strings.TrimSpace(projectID); pid != "" { + return filepath.Join(base, "projects", sanitizeEinoPathSegment(pid)) + } + conv := strings.TrimSpace(conversationID) + if conv == "" { + conv = "default" + } + return filepath.Join(base, "conversations", sanitizeEinoPathSegment(conv)) +} + +func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddlewareConfig, projectID, convID string, loc *localbk.Local, logger *zap.Logger) (adk.ChatModelAgentMiddleware, error) { + if loc == nil { + return nil, fmt.Errorf("reduction: local backend nil") + } + root := reductionCacheRootDir(mw.ReductionRootDir, projectID, convID) + if err := os.MkdirAll(root, 0o755); err != nil { + return nil, fmt.Errorf("reduction root: %w", err) + } + excl := append([]string(nil), mw.ReductionClearExclude...) + defaultExcl := []string{ + "task", "transfer_to_agent", "exit", "write_todos", "skill", "tool_search", + "TaskCreate", "TaskGet", "TaskUpdate", "TaskList", + } + excl = append(excl, defaultExcl...) + redMW, err := reduction.New(ctx, &reduction.Config{ + Backend: loc, + RootDir: root, + ReadFileToolName: "read_file", + ClearExcludeTools: excl, + MaxLengthForTrunc: mw.ReductionMaxLengthForTruncEffective(), + MaxTokensForClear: int64(mw.ReductionMaxTokensForClearEffective()), + }) + if err != nil { + return nil, err + } + if logger != nil { + logger.Info("eino middleware: reduction enabled", zap.String("root", root)) + } + return redMW, nil +} + +func buildAgenticReductionMiddleware( + ctx context.Context, + mw config.MultiAgentEinoMiddlewareConfig, + projectID, convID string, + loc *localbk.Local, + logger *zap.Logger, +) (adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], error) { + if loc == nil { + return nil, fmt.Errorf("agentic reduction: local backend nil") + } + root := reductionCacheRootDir(mw.ReductionRootDir, projectID, convID) + if err := os.MkdirAll(root, 0o755); err != nil { + return nil, fmt.Errorf("agentic reduction root: %w", err) + } + excl := append([]string(nil), mw.ReductionClearExclude...) + defaultExcl := []string{ + "task", "transfer_to_agent", "exit", "write_todos", "skill", "tool_search", + "TaskCreate", "TaskGet", "TaskUpdate", "TaskList", + } + excl = append(excl, defaultExcl...) + redMW, err := reduction.NewTyped[*schema.AgenticMessage](ctx, &reduction.TypedConfig[*schema.AgenticMessage]{ + Backend: loc, + RootDir: root, + ReadFileToolName: "read_file", + ClearExcludeTools: excl, + MaxLengthForTrunc: mw.ReductionMaxLengthForTruncEffective(), + MaxTokensForClear: int64(mw.ReductionMaxTokensForClearEffective()), + }) + if err != nil { + return nil, err + } + if logger != nil { + logger.Info("eino middleware: agentic reduction enabled", zap.String("root", root)) + } + return redMW, nil +} + +// prependEinoMiddlewares returns handlers to prepend (outermost first) and optionally replaces tools when tool_search is used. +// toolSearchActive is true when the toolsearch middleware was mounted (dynamic tools split off); callers should pass this to +// injectToolNamesOnlyInstruction — tool_search is not part of the pre-middleware tools list, so name-scanning alone cannot detect it. +func prependEinoMiddlewares( + ctx context.Context, + mw *config.MultiAgentEinoMiddlewareConfig, + place einoMWPlacement, + tools []tool.BaseTool, + einoLoc *localbk.Local, + skillsRoot string, + conversationID string, + projectID string, + logger *zap.Logger, +) (outTools []tool.BaseTool, extraHandlers []adk.ChatModelAgentMiddleware, toolSearchActive bool, err error) { + if mw == nil { + return tools, nil, false, nil + } + outTools = tools + + if mw.PatchToolCallsEffective() { + patchMW, perr := patchtoolcalls.New(ctx, &patchtoolcalls.Config{}) + if perr != nil { + return nil, nil, false, fmt.Errorf("patchtoolcalls: %w", perr) + } + extraHandlers = append(extraHandlers, patchMW) + } + + if mw.ReductionEnable && einoLoc != nil { + if place == einoMWSub && !mw.ReductionSubAgents { + // skip + } else { + redMW, rerr := buildReductionMiddleware(ctx, *mw, projectID, conversationID, einoLoc, logger) + if rerr != nil { + return nil, nil, false, rerr + } + extraHandlers = append(extraHandlers, redMW) + } + } + + minTools := mw.ToolSearchMinTools + if minTools <= 0 { + minTools = 20 + } + alwaysVis := mw.ToolSearchAlwaysVisible + if alwaysVis <= 0 { + alwaysVis = 12 + } + if mw.ToolSearchEnable && len(tools) >= minTools { + static, dynamic, split := splitToolsForToolSearchByNames(tools, mergeAlwaysVisibleToolNames(mw.ToolSearchAlwaysVisibleTools), alwaysVis) + if split && len(dynamic) > 0 { + ts, terr := toolsearch.New(ctx, &toolsearch.Config{DynamicTools: dynamic}) + if terr != nil { + return nil, nil, false, fmt.Errorf("toolsearch: %w", terr) + } + extraHandlers = append(extraHandlers, ts) + outTools = static + toolSearchActive = true + if logger != nil { + logger.Info("eino middleware: tool_search enabled", + zap.Int("static_tools", len(static)), + zap.Int("dynamic_tools", len(dynamic))) + } + } + } + + if place == einoMWMain && mw.PlantaskEnable { + if einoLoc == nil || strings.TrimSpace(skillsRoot) == "" { + if logger != nil { + logger.Warn("eino middleware: plantask_enable ignored (need eino_skills + skills_dir)") + } + } else { + rel := strings.TrimSpace(mw.PlantaskRelDir) + if rel == "" { + rel = ".eino/plantask" + } + baseDir := filepath.Join(skillsRoot, rel, sanitizeEinoPathSegment(conversationID)) + if mk := os.MkdirAll(baseDir, 0o755); mk != nil { + return nil, nil, toolSearchActive, fmt.Errorf("plantask mkdir: %w", mk) + } + ptBE := newLocalPlantaskBackend(einoLoc) + pt, perr := plantask.New(ctx, &plantask.Config{Backend: ptBE, BaseDir: baseDir}) + if perr != nil { + return nil, nil, toolSearchActive, fmt.Errorf("plantask: %w", perr) + } + extraHandlers = append(extraHandlers, pt) + if logger != nil { + logger.Info("eino middleware: plantask enabled", zap.String("baseDir", baseDir)) + } + } + } + + return outTools, extraHandlers, toolSearchActive, nil +} + +func prependEinoAgenticMiddlewares( + ctx context.Context, + mw *config.MultiAgentEinoMiddlewareConfig, + place einoMWPlacement, + tools []tool.BaseTool, + einoLoc *localbk.Local, + skillsRoot string, + conversationID string, + projectID string, + logger *zap.Logger, +) (outTools []tool.BaseTool, extraHandlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], toolSearchActive bool, err error) { + if mw == nil { + return tools, nil, false, nil + } + outTools = tools + + if mw.PatchToolCallsEffective() { + patchMW, perr := patchtoolcalls.NewTyped[*schema.AgenticMessage](ctx, &patchtoolcalls.Config{}) + if perr != nil { + return nil, nil, false, fmt.Errorf("agentic patchtoolcalls: %w", perr) + } + extraHandlers = append(extraHandlers, patchMW) + } + + if mw.ReductionEnable && einoLoc != nil { + if place == einoMWSub && !mw.ReductionSubAgents { + // skip + } else { + redMW, rerr := buildAgenticReductionMiddleware(ctx, *mw, projectID, conversationID, einoLoc, logger) + if rerr != nil { + return nil, nil, false, rerr + } + extraHandlers = append(extraHandlers, redMW) + } + } + + minTools := mw.ToolSearchMinTools + if minTools <= 0 { + minTools = 20 + } + alwaysVis := mw.ToolSearchAlwaysVisible + if alwaysVis <= 0 { + alwaysVis = 12 + } + if mw.ToolSearchEnable && len(tools) >= minTools { + static, dynamic, split := splitToolsForToolSearchByNames(tools, mergeAlwaysVisibleToolNames(mw.ToolSearchAlwaysVisibleTools), alwaysVis) + if split && len(dynamic) > 0 { + ts, terr := toolsearch.NewTyped[*schema.AgenticMessage](ctx, &toolsearch.Config{DynamicTools: dynamic}) + if terr != nil { + return nil, nil, false, fmt.Errorf("agentic toolsearch: %w", terr) + } + extraHandlers = append(extraHandlers, ts) + outTools = static + toolSearchActive = true + if logger != nil { + logger.Info("eino middleware: agentic tool_search enabled", + zap.Int("static_tools", len(static)), + zap.Int("dynamic_tools", len(dynamic))) + } + } + } + + if place == einoMWMain && mw.PlantaskEnable { + if einoLoc == nil || strings.TrimSpace(skillsRoot) == "" { + if logger != nil { + logger.Warn("eino middleware: agentic plantask_enable ignored (need eino_skills + skills_dir)") + } + } else { + rel := strings.TrimSpace(mw.PlantaskRelDir) + if rel == "" { + rel = ".eino/plantask" + } + baseDir := filepath.Join(skillsRoot, rel, sanitizeEinoPathSegment(conversationID)) + if mk := os.MkdirAll(baseDir, 0o755); mk != nil { + return nil, nil, toolSearchActive, fmt.Errorf("agentic plantask mkdir: %w", mk) + } + ptBE := newLocalPlantaskBackend(einoLoc) + pt, perr := plantask.NewTyped[*schema.AgenticMessage](ctx, &plantask.Config{Backend: ptBE, BaseDir: baseDir}) + if perr != nil { + return nil, nil, toolSearchActive, fmt.Errorf("agentic plantask: %w", perr) + } + extraHandlers = append(extraHandlers, pt) + if logger != nil { + logger.Info("eino middleware: agentic plantask enabled", zap.String("baseDir", baseDir)) + } + } + } + + return outTools, extraHandlers, toolSearchActive, nil +} + +func deepExtrasFromConfig(ma *config.MultiAgentConfig) (outputKey string, taskDesc func(context.Context, []adk.Agent) (string, error)) { + if ma == nil { + return "", nil + } + mw := ma.EinoMiddleware + if k := strings.TrimSpace(mw.DeepOutputKey); k != "" { + outputKey = k + } + prefix := strings.TrimSpace(mw.TaskToolDescriptionPrefix) + if prefix != "" { + taskDesc = func(ctx context.Context, agents []adk.Agent) (string, error) { + _ = ctx + var names []string + for _, a := range agents { + if a == nil { + continue + } + n := strings.TrimSpace(a.Name(ctx)) + if n != "" { + names = append(names, n) + } + } + if len(names) == 0 { + return prefix, nil + } + return prefix + "\n可用子代理(按名称 transfer / task 调用):" + strings.Join(names, "、"), nil + } + } + return outputKey, taskDesc +} + +func deepAgenticExtrasFromConfig(ma *config.MultiAgentConfig) (outputKey string, taskDesc func(context.Context, []adk.TypedAgent[*schema.AgenticMessage]) (string, error)) { + if ma == nil { + return "", nil + } + mw := ma.EinoMiddleware + if k := strings.TrimSpace(mw.DeepOutputKey); k != "" { + outputKey = k + } + prefix := strings.TrimSpace(mw.TaskToolDescriptionPrefix) + if prefix != "" { + taskDesc = func(ctx context.Context, agents []adk.TypedAgent[*schema.AgenticMessage]) (string, error) { + _ = ctx + var names []string + for _, a := range agents { + if a == nil { + continue + } + n := strings.TrimSpace(a.Name(ctx)) + if n != "" { + names = append(names, n) + } + } + if len(names) == 0 { + return prefix, nil + } + return prefix + "\n可用子代理(按名称 transfer / task 调用):" + strings.Join(names, "、"), nil + } + } + return outputKey, taskDesc +} diff --git a/internal/multiagent/eino_middleware_test.go b/internal/multiagent/eino_middleware_test.go new file mode 100644 index 00000000..45842d86 --- /dev/null +++ b/internal/multiagent/eino_middleware_test.go @@ -0,0 +1,220 @@ +package multiagent + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + localbk "github.com/cloudwego/eino-ext/adk/backend/local" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func TestReductionCacheRootDir(t *testing.T) { + got := reductionCacheRootDir("", "proj-1", "conv-1") + want := filepath.Join("tmp", "reduction", "projects", "proj-1") + if got != want { + t.Fatalf("project scope: got %q want %q", got, want) + } + got = reductionCacheRootDir("", "", "conv-abc") + want = filepath.Join("tmp", "reduction", "conversations", "conv-abc") + if got != want { + t.Fatalf("conversation scope: got %q want %q", got, want) + } + custom := reductionCacheRootDir("/data/cache", "p1", "c1") + if !strings.HasSuffix(custom, filepath.Join("projects", "p1")) { + t.Fatalf("custom base should still scope by project, got %q", custom) + } +} + +func TestBuildAgenticReductionMiddlewareClearsOldAgenticToolResult(t *testing.T) { + ctx := context.Background() + loc, err := localbk.NewBackend(ctx, &localbk.Config{}) + if err != nil { + t.Fatalf("NewBackend: %v", err) + } + root := t.TempDir() + mw, err := buildAgenticReductionMiddleware(ctx, config.MultiAgentEinoMiddlewareConfig{ + ReductionRootDir: root, + ReductionMaxTokensForClear: 1, + }, "", "conv-1", loc, nil) + if err != nil { + t.Fatalf("buildAgenticReductionMiddleware: %v", err) + } + oldText := strings.Repeat("old-tool-output-", 20) + newText := strings.Repeat("new-tool-output-", 20) + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + agenticAssistantToolCall("old-call", "execute", `{"command":"old"}`), + agenticToolResult("old-call", "execute", oldText), + agenticAssistantToolCall("new-call", "execute", `{"command":"new"}`), + agenticToolResult("new-call", "execute", newText), + }, + } + _, out, err := mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + oldGot := out.Messages[1].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text + newGot := out.Messages[3].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text + if oldGot == oldText { + t.Fatal("agentic reduction did not clear old oversized tool result") + } + if !strings.Contains(oldGot, "read_file") { + t.Fatalf("cleared content should mention read_file, got %q", oldGot) + } + if newGot != newText { + t.Fatalf("latest tool result should be retained, got %q", newGot) + } +} + +func agenticAssistantToolCall(callID, name, arguments string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: callID, + Name: name, + Arguments: arguments, + })}, + } +} + +func agenticToolResult(callID, name, text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: name, + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: text}, + }}, + })}, + } +} + +func TestBuildAgenticReductionMiddlewareHandlesSingleAgenticToolResult(t *testing.T) { + ctx := context.Background() + loc, err := localbk.NewBackend(ctx, &localbk.Config{}) + if err != nil { + t.Fatalf("NewBackend: %v", err) + } + mw, err := buildAgenticReductionMiddleware(ctx, config.MultiAgentEinoMiddlewareConfig{ + ReductionRootDir: t.TempDir(), + ReductionMaxTokensForClear: 1, + }, "", "conv-1", loc, nil) + if err != nil { + t.Fatalf("buildAgenticReductionMiddleware: %v", err) + } + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "call-1", + Name: "execute", + Content: []*schema.FunctionToolResultContentBlock{{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: strings.Repeat("tool-output-", 20)}, + }}, + })}, + }, + }, + } + _, out, err := mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + got := out.Messages[0].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text + if got != strings.Repeat("tool-output-", 20) { + t.Fatalf("single retained tool result should not be cleared, got %q", got) + } +} + +func TestPrependEinoAgenticMiddlewaresRespectsReductionPlacement(t *testing.T) { + ctx := context.Background() + loc, err := localbk.NewBackend(ctx, &localbk.Config{}) + if err != nil { + t.Fatalf("NewBackend: %v", err) + } + patchToolCalls := false + mw := &config.MultiAgentEinoMiddlewareConfig{ + ReductionEnable: true, + ReductionRootDir: t.TempDir(), + ReductionMaxTokensForClear: 100, + PatchToolCalls: &patchToolCalls, + } + _, mainHandlers, _, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWMain, nil, loc, "", "conv-1", "", nil) + if err != nil { + t.Fatalf("prepend main: %v", err) + } + if len(mainHandlers) != 1 { + t.Fatalf("main handlers = %d, want reduction", len(mainHandlers)) + } + _, subHandlers, _, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWSub, nil, loc, "", "conv-1", "", nil) + if err != nil { + t.Fatalf("prepend sub: %v", err) + } + if len(subHandlers) != 0 { + t.Fatalf("sub handlers = %d, want skipped when reduction_sub_agents=false", len(subHandlers)) + } + mw.ReductionSubAgents = true + _, subHandlers, _, err = prependEinoAgenticMiddlewares(ctx, mw, einoMWSub, nil, loc, "", "conv-1", "", nil) + if err != nil { + t.Fatalf("prepend sub enabled: %v", err) + } + if len(subHandlers) != 1 { + t.Fatalf("sub handlers = %d, want reduction when reduction_sub_agents=true", len(subHandlers)) + } +} + +func TestPrependEinoAgenticMiddlewaresMountsToolSearchAndPatchToolCalls(t *testing.T) { + ctx := context.Background() + mw := &config.MultiAgentEinoMiddlewareConfig{ + ToolSearchEnable: true, + ToolSearchMinTools: 20, + ToolSearchAlwaysVisible: 5, + } + outTools, handlers, toolSearchActive, err := prependEinoAgenticMiddlewares(ctx, mw, einoMWMain, stubTools(25), nil, "", "conv-test", "", nil) + if err != nil { + t.Fatalf("prependEinoAgenticMiddlewares: %v", err) + } + if !toolSearchActive { + t.Fatal("agentic tool_search should be active") + } + if len(outTools) != 5 { + t.Fatalf("mounted tools = %d, want static visible tools only", len(outTools)) + } + if len(handlers) != 2 { + t.Fatalf("handlers = %d, want patchtoolcalls + toolsearch", len(handlers)) + } +} + +type stubTool struct{ name string } + +func (s stubTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: s.name}, nil +} + +func TestSplitToolsForToolSearch(t *testing.T) { + mk := func(n int) []tool.BaseTool { + out := make([]tool.BaseTool, n) + for i := 0; i < n; i++ { + out[i] = stubTool{name: fmt.Sprintf("t%d", i)} + } + return out + } + static, dynamic, ok := splitToolsForToolSearch(mk(4), 3) + if ok || len(static) != 4 || dynamic != nil { + t.Fatalf("expected no split when len<=alwaysVisible+1, got ok=%v static=%d dynamic=%v", ok, len(static), dynamic) + } + static, dynamic, ok = splitToolsForToolSearch(mk(20), 5) + if !ok || len(static) != 5 || len(dynamic) != 15 { + t.Fatalf("expected split 5+15, got ok=%v static=%d dynamic=%d", ok, len(static), len(dynamic)) + } +} diff --git a/internal/multiagent/eino_model_facing_persistence_test.go b/internal/multiagent/eino_model_facing_persistence_test.go new file mode 100644 index 00000000..eddf1bed --- /dev/null +++ b/internal/multiagent/eino_model_facing_persistence_test.go @@ -0,0 +1,34 @@ +package multiagent + +import ( + "strings" + "testing" + + "cyberstrike-ai/internal/agent" + + "github.com/cloudwego/eino/schema" +) + +func TestBuildEinoRunResultNeverPersistsRawAccumulationWithoutModelFacingTrace(t *testing.T) { + raw := []schema.Message{*schema.ToolMessage(strings.Repeat("raw-tool-output", 1000), "call-1")} + rawMsgs := make([]*schema.Message, len(raw)) + for i := range raw { + rawMsgs[i] = &raw[i] + } + result := buildEinoRunResultFromAccumulated("deep", rawMsgs, nil, "", "", "empty", nil, true) + if result.LastAgentTraceInput != "" { + t.Fatalf("pre-model raw accumulation must not be persisted: %d bytes", len(result.LastAgentTraceInput)) + } + + modelFacing := []*schema.Message{schema.UserMessage("bounded-model-view")} + result = buildEinoRunResultFromAccumulated("deep", rawMsgs, modelFacing, "ok", "", "empty", nil, false) + if !strings.Contains(result.LastAgentTraceInput, "bounded-model-view") { + t.Fatalf("model-facing trace missing: %s", result.LastAgentTraceInput) + } + if strings.Contains(result.LastAgentTraceInput, "raw-tool-output") { + t.Fatal("raw accumulation leaked into persisted model-facing trace") + } + if !agent.IsModelFacingTraceJSON(result.LastAgentTraceInput) { + t.Fatal("persisted model-facing trace is missing its version marker") + } +} diff --git a/internal/multiagent/eino_model_facing_trace.go b/internal/multiagent/eino_model_facing_trace.go new file mode 100644 index 00000000..33d8d011 --- /dev/null +++ b/internal/multiagent/eino_model_facing_trace.go @@ -0,0 +1,124 @@ +package multiagent + +import ( + "context" + "encoding/json" + "sync" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// modelFacingTraceHolder 保存「即将送入 ChatModel」的消息快照(已走 summarization / reduction / orphan 修剪等), +// 用于 last_react_input 落库,使续跑与「上下文压缩后」的模型视角一致,而非仅依赖事件流 append 的 runAccumulatedMsgs。 +type modelFacingTraceHolder struct { + mu sync.Mutex + // msgs 为深拷贝后的切片,避免框架后续原地修改污染快照 + msgs []adk.Message +} + +func newModelFacingTraceHolder() *modelFacingTraceHolder { + return &modelFacingTraceHolder{} +} + +// Snapshot 返回当前快照的再一次深拷贝(供序列化落库,避免与 holder 互斥长期持锁)。 +func (h *modelFacingTraceHolder) Snapshot() []adk.Message { + if h == nil { + return nil + } + h.mu.Lock() + defer h.mu.Unlock() + return cloneADKMessagesForTrace(h.msgs) +} + +func (h *modelFacingTraceHolder) storeFromState(state *adk.ChatModelAgentState) { + if h == nil || state == nil || len(state.Messages) == 0 { + return + } + cloned := cloneADKMessagesForTrace(state.Messages) + if len(cloned) == 0 { + return + } + h.mu.Lock() + h.msgs = cloned + h.mu.Unlock() +} + +func (h *modelFacingTraceHolder) storeFromAgenticState(state *adk.TypedChatModelAgentState[*schema.AgenticMessage]) { + if h == nil || state == nil || len(state.Messages) == 0 { + return + } + cloned := cloneADKMessagesForTrace(AgenticMessagesToEino(state.Messages)) + if len(cloned) == 0 { + return + } + h.mu.Lock() + h.msgs = cloned + h.mu.Unlock() +} + +func cloneADKMessagesForTrace(msgs []adk.Message) []adk.Message { + if len(msgs) == 0 { + return nil + } + b, err := json.Marshal(msgs) + if err != nil { + return nil + } + var out []adk.Message + if err := json.Unmarshal(b, &out); err != nil { + return nil + } + return out +} + +// modelFacingTraceMiddleware 必须在 Handlers 链中处于 **BeforeModel 最后**(telemetry 之后), +// 此时 state.Messages 即为本次 LLM 调用的最终入参。 +type modelFacingTraceMiddleware struct { + adk.BaseChatModelAgentMiddleware + holder *modelFacingTraceHolder +} + +func newModelFacingTraceMiddleware(holder *modelFacingTraceHolder) adk.ChatModelAgentMiddleware { + if holder == nil { + return nil + } + return &modelFacingTraceMiddleware{holder: holder} +} + +func (m *modelFacingTraceMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.ChatModelAgentState, + mc *adk.ModelContext, +) (context.Context, *adk.ChatModelAgentState, error) { + if m.holder != nil && state != nil { + m.holder.storeFromState(state) + } + return ctx, state, nil +} + +type agenticModelFacingTraceMiddleware struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + holder *modelFacingTraceHolder +} + +func newAgenticModelFacingTraceMiddleware(holder *modelFacingTraceHolder) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + if holder == nil { + return nil + } + return &agenticModelFacingTraceMiddleware{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + holder: holder, + } +} + +func (m *agenticModelFacingTraceMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + mc *adk.TypedModelContext[*schema.AgenticMessage], +) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m.holder != nil && state != nil { + m.holder.storeFromAgenticState(state) + } + return ctx, state, nil +} diff --git a/internal/multiagent/eino_model_resilience.go b/internal/multiagent/eino_model_resilience.go new file mode 100644 index 00000000..4725a339 --- /dev/null +++ b/internal/multiagent/eino_model_resilience.go @@ -0,0 +1,538 @@ +package multiagent + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/llm" + "cyberstrike-ai/internal/openai" + "cyberstrike-ai/internal/reasoning" + + agenticopenai "github.com/cloudwego/eino-ext/components/model/agenticopenai" + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +type einoModelMode string + +const ( + einoModelModeNormal einoModelMode = "normal" + einoModelModePlanner einoModelMode = "planner" +) + +type einoModelFactory func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.ToolCallingChatModel, error) +type einoAgenticModelConfigFactory func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.AgenticModel, error) + +func newEinoBaseHTTPClient() *http.Client { + return &http.Client{ + Timeout: 30 * time.Minute, + Transport: &http.Transport{ + DialContext: (&net.Dialer{ + Timeout: 300 * time.Second, + KeepAlive: 300 * time.Second, + }).DialContext, + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 30 * time.Second, + ResponseHeaderTimeout: 60 * time.Minute, + }, + } +} + +func newEinoToolCallingChatModelFactory( + baseHTTPClient *http.Client, + reasoningClient *reasoning.ClientIntent, + logger *zap.Logger, +) einoModelFactory { + if baseHTTPClient == nil { + baseHTTPClient = newEinoBaseHTTPClient() + } + return func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.ToolCallingChatModel, error) { + if isEinoAgenticClaudeProvider(oa.Provider) { + nativeModel, err := newEinoClaudeAgenticChatModel(ctx, oa, mode, baseHTTPClient, reasoningClient) + if err != nil { + return nil, err + } + return newAgenticToolCallingChatModelAdapter(nativeModel), nil + } + httpClient := openai.NewEinoHTTPClient(&oa, baseHTTPClient) + openai.AttachSummarizationDiagTransport(httpClient, logger) + maxCompletionTokens := oa.MaxCompletionTokensEffective() + modelCfg := &einoopenai.ChatModelConfig{ + APIKey: oa.APIKey, + BaseURL: strings.TrimSuffix(oa.BaseURL, "/"), + Model: oa.Model, + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + } + if mode == einoModelModePlanner { + reasoning.ApplyPlanExecutePlannerModelConfig(modelCfg, &oa) + } else { + reasoning.ApplyToEinoChatModelConfig(modelCfg, &oa, reasoningClient) + } + baseModel, err := einoopenai.NewChatModel(ctx, modelCfg) + if err != nil { + return nil, err + } + return newStreamToolCallIndexRepairModel(baseModel), nil + } +} + +func newEinoAgenticChatModelFactory( + baseHTTPClient *http.Client, + reasoningClient *reasoning.ClientIntent, + logger *zap.Logger, +) einoAgenticModelConfigFactory { + if baseHTTPClient == nil { + baseHTTPClient = newEinoBaseHTTPClient() + } + return func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.AgenticModel, error) { + if !supportsEinoAgenticBackend(oa) { + 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) + } + httpClient := openai.NewEinoHTTPClient(&oa, baseHTTPClient) + openai.AttachSummarizationDiagTransport(httpClient, logger) + maxCompletionTokens := oa.MaxCompletionTokensEffective() + modelCfg := &agenticopenai.ChatConfig{ + APIKey: oa.APIKey, + BaseURL: strings.TrimSuffix(oa.BaseURL, "/"), + Model: oa.Model, + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + ExtraFields: reasoning.AgenticOpenAIExtraFields(&oa, reasoningClient), + } + if mode == einoModelModePlanner { + modelCfg.ExtraFields = reasoning.AgenticOpenAIPlannerExtraFields(&oa) + } + return agenticopenai.NewChatModel(ctx, modelCfg) + } +} + +func newEinoClaudeAgenticChatModel( + ctx context.Context, + oa config.OpenAIConfig, + mode einoModelMode, + httpClient *http.Client, + reasoningClient *reasoning.ClientIntent, +) (model.AgenticModel, error) { + extraFields := reasoning.AgenticOpenAIExtraFields(&oa, reasoningClient) + if mode == einoModelModePlanner { + extraFields = reasoning.AgenticOpenAIPlannerExtraFields(&oa) + } + return llm.NewClaudeAgenticModel( + ctx, + oa, + httpClient, + oa.MaxCompletionTokensEffective(), + extraFields, + ) +} + +func supportsEinoAgenticBackend(oa config.OpenAIConfig) bool { + provider := strings.ToLower(strings.TrimSpace(oa.Provider)) + return provider == "" || + provider == "openai" || + provider == "openai_compatible" || + isEinoAgenticClaudeProvider(provider) +} + +func isEinoAgenticClaudeProvider(provider string) bool { + return llm.IsClaudeProvider(provider) +} + +func agenticModelGateFactory(factory einoAgenticModelConfigFactory, oa config.OpenAIConfig, mode einoModelMode) einoAgenticModelFactory { + if factory == nil { + return nil + } + return func(ctx context.Context) (model.AgenticModel, error) { + return factory(ctx, oa, mode) + } +} + +func newEinoModelRetryConfig( + mw *config.MultiAgentEinoMiddlewareConfig, + logger *zap.Logger, + scope string, +) *adk.ModelRetryConfig { + maxRetries := RunRetryMaxAttemptsFromConfig(mw) + maxBackoff := einoRunRetryMaxBackoffFromConfig(mw) + return &adk.ModelRetryConfig{ + MaxRetries: maxRetries, + BackoffFunc: func(_ context.Context, attempt int) time.Duration { + return einoTransientRetryBackoff(attempt-1, maxBackoff) + }, + ShouldRetry: func(ctx context.Context, retryCtx *adk.RetryContext) *adk.RetryDecision { + if retryCtx == nil || ctx.Err() != nil { + return &adk.RetryDecision{} + } + if retryCtx.Err != nil { + if !isEinoTransientRunError(retryCtx.Err) { + return &adk.RetryDecision{} + } + if logger != nil { + kind, summary := einoTransientRunErrorUserDetail(retryCtx.Err) + logger.Warn("eino native model retry", + zap.String("scope", scope), + zap.Int("attempt", retryCtx.RetryAttempt), + zap.Int("maxRetries", maxRetries), + zap.String("errorKind", kind), + zap.String("errorSummary", summary), + ) + } + return &adk.RetryDecision{Retry: true, RejectReason: "transient_model_error"} + } + if isRetryableEmptyModelOutput(retryCtx.OutputMessage) { + if logger != nil { + logger.Warn("eino native model retry: empty model output", + zap.String("scope", scope), + zap.Int("attempt", retryCtx.RetryAttempt), + zap.Int("maxRetries", maxRetries), + ) + } + return &adk.RetryDecision{Retry: true, RejectReason: "empty_model_output"} + } + return &adk.RetryDecision{} + }, + } +} + +func newEinoAgenticModelRetryConfig( + mw *config.MultiAgentEinoMiddlewareConfig, + logger *zap.Logger, + scope string, +) *adk.TypedModelRetryConfig[*schema.AgenticMessage] { + maxRetries := RunRetryMaxAttemptsFromConfig(mw) + maxBackoff := einoRunRetryMaxBackoffFromConfig(mw) + return &adk.TypedModelRetryConfig[*schema.AgenticMessage]{ + MaxRetries: maxRetries, + BackoffFunc: func(_ context.Context, attempt int) time.Duration { + return einoTransientRetryBackoff(attempt-1, maxBackoff) + }, + ShouldRetry: func(ctx context.Context, retryCtx *adk.TypedRetryContext[*schema.AgenticMessage]) *adk.TypedRetryDecision[*schema.AgenticMessage] { + if retryCtx == nil || ctx.Err() != nil { + return &adk.TypedRetryDecision[*schema.AgenticMessage]{} + } + if retryCtx.Err != nil { + if !isEinoTransientRunError(retryCtx.Err) { + return &adk.TypedRetryDecision[*schema.AgenticMessage]{} + } + if logger != nil { + kind, summary := einoTransientRunErrorUserDetail(retryCtx.Err) + logger.Warn("eino native agentic model retry", + zap.String("scope", scope), + zap.Int("attempt", retryCtx.RetryAttempt), + zap.Int("maxRetries", maxRetries), + zap.String("errorKind", kind), + zap.String("errorSummary", summary), + ) + } + return &adk.TypedRetryDecision[*schema.AgenticMessage]{Retry: true, RejectReason: "transient_model_error"} + } + if isRetryableEmptyAgenticModelOutput(retryCtx.OutputMessage) { + if logger != nil { + logger.Warn("eino native agentic model retry: empty model output", + zap.String("scope", scope), + zap.Int("attempt", retryCtx.RetryAttempt), + zap.Int("maxRetries", maxRetries), + ) + } + return &adk.TypedRetryDecision[*schema.AgenticMessage]{Retry: true, RejectReason: "empty_model_output"} + } + return &adk.TypedRetryDecision[*schema.AgenticMessage]{} + }, + } +} + +func newEinoModelFailoverConfig( + ctx context.Context, + appCfg *config.Config, + mw *config.MultiAgentEinoMiddlewareConfig, + mode einoModelMode, + factory einoModelFactory, + logger *zap.Logger, + scope string, + progress func(eventType, message string, data interface{}), + orchestration string, + conversationID string, +) (*adk.ModelFailoverConfig[*schema.Message], error) { + channels := resolveEinoFailoverChannels(appCfg, mw) + if len(channels) == 0 { + return nil, nil + } + if factory == nil { + return nil, fmt.Errorf("eino model failover: 模型工厂为空") + } + + maxRetries := len(channels) + if mw != nil && mw.ModelFailoverMaxRetries > 0 && mw.ModelFailoverMaxRetries < maxRetries { + maxRetries = mw.ModelFailoverMaxRetries + } + channels = channels[:maxRetries] + + cache := make(map[string]model.BaseModel[*schema.Message], len(channels)) + var mu sync.Mutex + return &adk.ModelFailoverConfig[*schema.Message]{ + MaxRetries: uint(maxRetries), + ShouldFailover: func(ctx context.Context, _ *schema.Message, err error) bool { + if ctx.Err() != nil || err == nil { + return false + } + err = unwrapEinoRetryExhausted(err) + return isEinoTransientRunError(err) + }, + GetFailoverModel: func(ctx context.Context, failoverCtx *adk.FailoverContext[*schema.Message]) (model.BaseModel[*schema.Message], []*schema.Message, error) { + if failoverCtx == nil || failoverCtx.FailoverAttempt == 0 { + return nil, nil, fmt.Errorf("eino model failover: invalid failover attempt") + } + idx := int(failoverCtx.FailoverAttempt) - 1 + if idx < 0 || idx >= len(channels) { + return nil, nil, fmt.Errorf("eino model failover: no channel for attempt %d", failoverCtx.FailoverAttempt) + } + ch := channels[idx] + mu.Lock() + cached := cache[ch.id] + mu.Unlock() + if cached != nil { + emitEinoModelFailoverEvent(progress, conversationID, orchestration, scope, ch.id, ch.cfg.Model, failoverCtx.FailoverAttempt) + if logger != nil { + logger.Warn("eino native model failover", + zap.String("scope", scope), + zap.String("channel", ch.id), + zap.String("model", ch.cfg.Model), + zap.Uint("attempt", failoverCtx.FailoverAttempt), + ) + } + return cached, nil, nil + } + m, err := factory(ctx, ch.cfg, mode) + if err != nil { + return nil, nil, fmt.Errorf("eino model failover channel %q: %w", ch.id, err) + } + mu.Lock() + cache[ch.id] = m + mu.Unlock() + emitEinoModelFailoverEvent(progress, conversationID, orchestration, scope, ch.id, ch.cfg.Model, failoverCtx.FailoverAttempt) + if logger != nil { + logger.Warn("eino native model failover", + zap.String("scope", scope), + zap.String("channel", ch.id), + zap.String("model", ch.cfg.Model), + zap.Uint("attempt", failoverCtx.FailoverAttempt), + ) + } + return m, nil, nil + }, + }, nil +} + +func newEinoAgenticModelFailoverConfig( + ctx context.Context, + appCfg *config.Config, + mw *config.MultiAgentEinoMiddlewareConfig, + mode einoModelMode, + factory einoAgenticModelConfigFactory, + logger *zap.Logger, + scope string, + progress func(eventType, message string, data interface{}), + orchestration string, + conversationID string, +) (*adk.ModelFailoverConfig[*schema.AgenticMessage], error) { + channels := resolveEinoFailoverChannels(appCfg, mw) + if len(channels) == 0 { + return nil, nil + } + if factory == nil { + return nil, fmt.Errorf("eino agentic model failover: 模型工厂为空") + } + + maxRetries := len(channels) + if mw != nil && mw.ModelFailoverMaxRetries > 0 && mw.ModelFailoverMaxRetries < maxRetries { + maxRetries = mw.ModelFailoverMaxRetries + } + channels = channels[:maxRetries] + + cache := make(map[string]model.BaseModel[*schema.AgenticMessage], len(channels)) + var mu sync.Mutex + return &adk.ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: uint(maxRetries), + ShouldFailover: func(ctx context.Context, _ *schema.AgenticMessage, err error) bool { + if ctx.Err() != nil || err == nil { + return false + } + err = unwrapEinoRetryExhausted(err) + return isEinoTransientRunError(err) + }, + GetFailoverModel: func(ctx context.Context, failoverCtx *adk.FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + if failoverCtx == nil || failoverCtx.FailoverAttempt == 0 { + return nil, nil, fmt.Errorf("eino agentic model failover: invalid failover attempt") + } + idx := int(failoverCtx.FailoverAttempt) - 1 + if idx < 0 || idx >= len(channels) { + return nil, nil, fmt.Errorf("eino agentic model failover: no channel for attempt %d", failoverCtx.FailoverAttempt) + } + ch := channels[idx] + mu.Lock() + cached := cache[ch.id] + mu.Unlock() + if cached != nil { + emitEinoModelFailoverEvent(progress, conversationID, orchestration, scope, ch.id, ch.cfg.Model, failoverCtx.FailoverAttempt) + if logger != nil { + logger.Warn("eino native agentic model failover", + zap.String("scope", scope), + zap.String("channel", ch.id), + zap.String("model", ch.cfg.Model), + zap.Uint("attempt", failoverCtx.FailoverAttempt), + ) + } + return cached, nil, nil + } + m, err := factory(ctx, ch.cfg, mode) + if err != nil { + return nil, nil, fmt.Errorf("eino agentic model failover channel %q: %w", ch.id, err) + } + mu.Lock() + cache[ch.id] = m + mu.Unlock() + emitEinoModelFailoverEvent(progress, conversationID, orchestration, scope, ch.id, ch.cfg.Model, failoverCtx.FailoverAttempt) + if logger != nil { + logger.Warn("eino native agentic model failover", + zap.String("scope", scope), + zap.String("channel", ch.id), + zap.String("model", ch.cfg.Model), + zap.Uint("attempt", failoverCtx.FailoverAttempt), + ) + } + return m, nil, nil + }, + }, nil +} + +type resolvedEinoFailoverChannel struct { + id string + cfg config.OpenAIConfig +} + +func resolveEinoFailoverChannels(appCfg *config.Config, mw *config.MultiAgentEinoMiddlewareConfig) []resolvedEinoFailoverChannel { + if appCfg == nil || mw == nil || len(mw.ModelFailoverChannels) == 0 { + return nil + } + primary := appCfg.OpenAI + seen := map[string]struct{}{} + out := make([]resolvedEinoFailoverChannel, 0, len(mw.ModelFailoverChannels)) + for _, raw := range mw.ModelFailoverChannels { + id := config.NormalizeAIChannelID(raw) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + oa, resolvedID, ok := appCfg.AI.ResolveChannel(id) + if !ok { + continue + } + if sameOpenAIModelEndpoint(primary, oa) { + continue + } + seen[resolvedID] = struct{}{} + out = append(out, resolvedEinoFailoverChannel{id: resolvedID, cfg: oa}) + } + return out +} + +func sameOpenAIModelEndpoint(a, b config.OpenAIConfig) bool { + return strings.EqualFold(strings.TrimSpace(a.Provider), strings.TrimSpace(b.Provider)) && + strings.TrimRight(strings.TrimSpace(a.BaseURL), "/") == strings.TrimRight(strings.TrimSpace(b.BaseURL), "/") && + strings.TrimSpace(a.APIKey) == strings.TrimSpace(b.APIKey) && + strings.TrimSpace(a.Model) == strings.TrimSpace(b.Model) +} + +func isRetryableEmptyModelOutput(msg *schema.Message) bool { + if msg == nil { + return true + } + return strings.TrimSpace(msg.Content) == "" && + strings.TrimSpace(msg.ReasoningContent) == "" && + len(msg.ToolCalls) == 0 && + len(msg.MultiContent) == 0 && + len(msg.UserInputMultiContent) == 0 && + len(msg.AssistantGenMultiContent) == 0 +} + +func isRetryableEmptyAgenticModelOutput(msg *schema.AgenticMessage) bool { + if msg == nil { + return true + } + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + switch { + case block.Reasoning != nil: + if strings.TrimSpace(block.Reasoning.Text) != "" { + return false + } + case block.UserInputText != nil: + if strings.TrimSpace(block.UserInputText.Text) != "" { + return false + } + case block.AssistantGenText != nil: + if strings.TrimSpace(block.AssistantGenText.Text) != "" { + return false + } + default: + return false + } + } + return true +} + +func unwrapEinoRetryExhausted(err error) error { + var retryErr *adk.RetryExhaustedError + if errors.As(err, &retryErr) && retryErr.LastErr != nil { + return retryErr.LastErr + } + return err +} + +func isEinoNativeWillRetry(err error) (*adk.WillRetryError, bool) { + var willRetry *adk.WillRetryError + if errors.As(err, &willRetry) { + return willRetry, true + } + return nil, false +} + +func emitEinoModelFailoverEvent( + progress func(eventType, message string, data interface{}), + conversationID, orchestration, scope, channelID, modelName string, + attempt uint, +) { + if progress == nil { + return + } + msg := fmt.Sprintf("主模型重试耗尽,正在切换备用模型 %s。", modelName) + progress("eino_model_failover", msg, map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "orchestration": orchestration, + "scope": scope, + "channel": channelID, + "model": modelName, + "attempt": attempt, + }) +} diff --git a/internal/multiagent/eino_model_resilience_test.go b/internal/multiagent/eino_model_resilience_test.go new file mode 100644 index 00000000..8f7f5108 --- /dev/null +++ b/internal/multiagent/eino_model_resilience_test.go @@ -0,0 +1,401 @@ +package multiagent + +import ( + "context" + "errors" + "testing" + "time" + + "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" +) + +func TestNewEinoModelRetryConfigUsesNativeFieldsFirst(t *testing.T) { + t.Parallel() + mw := &config.MultiAgentEinoMiddlewareConfig{ + ModelRetryMaxRetries: 2, + ModelRetryMaxBackoffSec: 7, + RunRetryMaxAttempts: 9, + RunRetryMaxBackoffSec: 11, + } + cfg := newEinoModelRetryConfig(mw, nil, "test") + if cfg.MaxRetries != 2 { + t.Fatalf("MaxRetries = %d, want 2", cfg.MaxRetries) + } + backoff := cfg.BackoffFunc(context.Background(), 1) + if backoff < 500*time.Millisecond || backoff > 2*time.Second { + t.Fatalf("attempt 1 backoff = %v, want first equal-jitter window", backoff) + } + if got := einoRunRetryMaxBackoffFromConfig(mw); got != 7*time.Second { + t.Fatalf("backoff from config = %v, want 7s", got) + } +} + +func TestEinoModelRetryPolicyRetriesTransientAndEmptyOutput(t *testing.T) { + t.Parallel() + cfg := newEinoModelRetryConfig(&config.MultiAgentEinoMiddlewareConfig{ModelRetryMaxRetries: 1}, nil, "test") + if got := cfg.ShouldRetry(context.Background(), &adk.RetryContext{Err: errors.New("HTTP 429 Too Many Requests")}); got == nil || !got.Retry { + t.Fatal("transient model error should retry") + } + if got := cfg.ShouldRetry(context.Background(), &adk.RetryContext{OutputMessage: schema.AssistantMessage("", nil)}); got == nil || !got.Retry { + t.Fatal("empty assistant output should retry") + } + if got := cfg.ShouldRetry(context.Background(), &adk.RetryContext{OutputMessage: schema.AssistantMessage("", []schema.ToolCall{{ID: "call_1"}})}); got == nil || got.Retry { + t.Fatal("assistant tool call output should not be treated as empty") + } + if got := cfg.ShouldRetry(context.Background(), &adk.RetryContext{Err: errors.New("invalid api key")}); got == nil || got.Retry { + t.Fatal("permanent auth error should not retry") + } +} + +func TestEinoAgenticModelRetryPolicyRetriesTransientAndEmptyOutput(t *testing.T) { + t.Parallel() + cfg := newEinoAgenticModelRetryConfig(&config.MultiAgentEinoMiddlewareConfig{ModelRetryMaxRetries: 1}, nil, "agentic") + if got := cfg.ShouldRetry(context.Background(), &adk.TypedRetryContext[*schema.AgenticMessage]{Err: errors.New("HTTP 429 Too Many Requests")}); got == nil || !got.Retry { + t.Fatal("transient agentic model error should retry") + } + if got := cfg.ShouldRetry(context.Background(), &adk.TypedRetryContext[*schema.AgenticMessage]{ + OutputMessage: &schema.AgenticMessage{Role: schema.AgenticRoleTypeAssistant}, + }); got == nil || !got.Retry { + t.Fatal("empty agentic assistant output should retry") + } + if got := cfg.ShouldRetry(context.Background(), &adk.TypedRetryContext[*schema.AgenticMessage]{ + OutputMessage: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: "ok"})}, + }, + }); got == nil || got.Retry { + t.Fatal("agentic assistant text should not be treated as empty") + } + if got := cfg.ShouldRetry(context.Background(), &adk.TypedRetryContext[*schema.AgenticMessage]{ + OutputMessage: &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call_1", Name: "search", Arguments: `{"q":"x"}`, + })}, + }, + }); got == nil || got.Retry { + t.Fatal("agentic tool call output should not be treated as empty") + } + if got := cfg.ShouldRetry(context.Background(), &adk.TypedRetryContext[*schema.AgenticMessage]{Err: errors.New("invalid api key")}); got == nil || got.Retry { + t.Fatal("permanent auth error should not retry") + } +} + +func TestResolveEinoFailoverChannelsSkipsPrimaryDuplicateAndUnknown(t *testing.T) { + t.Parallel() + appCfg := &config.Config{ + OpenAI: config.OpenAIConfig{Provider: "openai", APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + AI: config.AIConfig{Channels: map[string]config.AIChannelConfig{ + "same": {Provider: "openai", APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + "fb1": {Provider: "openai", APIKey: "k2", BaseURL: "https://api.example/v1", Model: "fallback-1"}, + "fb2": {Provider: "claude", APIKey: "k3", BaseURL: "https://api.anthropic.com/v1", Model: "claude-sonnet"}, + }}, + } + got := resolveEinoFailoverChannels(appCfg, &config.MultiAgentEinoMiddlewareConfig{ + ModelFailoverChannels: []string{"same", "missing", "fb1", "fb1", "fb2"}, + ModelFailoverMaxRetries: 1, + }) + if len(got) != 2 { + t.Fatalf("resolved channels len = %d, want 2 before max cap is applied by config builder", len(got)) + } + if got[0].id != "fb1" || got[1].id != "fb2" { + t.Fatalf("resolved channel order = %#v", got) + } +} + +func TestNewEinoModelFailoverConfigBuildsDistinctFallbackModel(t *testing.T) { + t.Parallel() + appCfg := &config.Config{ + OpenAI: config.OpenAIConfig{APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + AI: config.AIConfig{Channels: map[string]config.AIChannelConfig{ + "fb1": {APIKey: "k2", BaseURL: "https://api.example/v1", Model: "fallback-1"}, + "fb2": {APIKey: "k3", BaseURL: "https://api.example/v1", Model: "fallback-2"}, + }}, + } + var built []string + cfg, err := newEinoModelFailoverConfig( + context.Background(), + appCfg, + &config.MultiAgentEinoMiddlewareConfig{ + ModelFailoverChannels: []string{"fb1", "fb2"}, + ModelFailoverMaxRetries: 1, + }, + einoModelModeNormal, + func(_ context.Context, oa config.OpenAIConfig, _ einoModelMode) (model.ToolCallingChatModel, error) { + built = append(built, oa.Model) + return &streamToolCallIndexFakeModel{}, nil + }, + nil, + "test", + nil, + "deep", + "conv-1", + ) + if err != nil { + t.Fatalf("newEinoModelFailoverConfig: %v", err) + } + if cfg == nil || cfg.MaxRetries != 1 { + t.Fatalf("failover cfg = %#v, want max retries 1", cfg) + } + m, msgs, err := cfg.GetFailoverModel(context.Background(), &adk.FailoverContext[*schema.Message]{FailoverAttempt: 1}) + if err != nil || m == nil || msgs != nil { + t.Fatalf("GetFailoverModel = (%v, %v, %v)", m, msgs, err) + } + if len(built) != 1 || built[0] != "fallback-1" { + t.Fatalf("built models = %v, want [fallback-1]", built) + } + if !cfg.ShouldFailover(context.Background(), nil, &adk.RetryExhaustedError{LastErr: errors.New("upstream returned 503"), TotalRetries: 4}) { + t.Fatal("retry-exhausted transient error should fail over") + } + if cfg.ShouldFailover(context.Background(), nil, &adk.RetryExhaustedError{LastErr: errors.New("invalid api key"), TotalRetries: 4}) { + t.Fatal("retry-exhausted permanent error should not fail over") + } +} + +func TestNewEinoModelFailoverConfigEmitsProgressEvent(t *testing.T) { + t.Parallel() + appCfg := &config.Config{ + OpenAI: config.OpenAIConfig{APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + AI: config.AIConfig{Channels: map[string]config.AIChannelConfig{ + "fb1": {APIKey: "k2", BaseURL: "https://api.example/v1", Model: "fallback-1"}, + }}, + } + var events []struct { + eventType string + message string + data interface{} + } + cfg, err := newEinoModelFailoverConfig( + context.Background(), + appCfg, + &config.MultiAgentEinoMiddlewareConfig{ModelFailoverChannels: []string{"fb1"}}, + einoModelModeNormal, + func(_ context.Context, _ config.OpenAIConfig, _ einoModelMode) (model.ToolCallingChatModel, error) { + return &streamToolCallIndexFakeModel{}, nil + }, + nil, + "test", + func(eventType, message string, data interface{}) { + events = append(events, struct { + eventType string + message string + data interface{} + }{eventType: eventType, message: message, data: data}) + }, + "deep", + "conv-1", + ) + if err != nil { + t.Fatalf("newEinoModelFailoverConfig: %v", err) + } + if _, _, err := cfg.GetFailoverModel(context.Background(), &adk.FailoverContext[*schema.Message]{FailoverAttempt: 1}); err != nil { + t.Fatalf("GetFailoverModel: %v", err) + } + if len(events) != 1 || events[0].eventType != "eino_model_failover" { + t.Fatalf("events = %#v, want one eino_model_failover", events) + } + payload, ok := events[0].data.(map[string]interface{}) + if !ok { + t.Fatalf("event payload type = %T", events[0].data) + } + if payload["conversationId"] != "conv-1" || payload["orchestration"] != "deep" || payload["channel"] != "fb1" || payload["model"] != "fallback-1" { + t.Fatalf("payload = %#v", payload) + } +} + +func TestNewEinoAgenticModelFailoverConfigBuildsDistinctFallbackModel(t *testing.T) { + t.Parallel() + appCfg := &config.Config{ + OpenAI: config.OpenAIConfig{Provider: "openai", APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + AI: config.AIConfig{Channels: map[string]config.AIChannelConfig{ + "fb1": {Provider: "openai", APIKey: "k2", BaseURL: "https://api.example/v1", Model: "fallback-1"}, + "fb2": {Provider: "openai", APIKey: "k3", BaseURL: "https://api.example/v1", Model: "fallback-2"}, + }}, + } + var built []string + cfg, err := newEinoAgenticModelFailoverConfig( + context.Background(), + appCfg, + &config.MultiAgentEinoMiddlewareConfig{ + ModelFailoverChannels: []string{"fb1", "fb2"}, + ModelFailoverMaxRetries: 1, + }, + einoModelModeNormal, + func(_ context.Context, oa config.OpenAIConfig, _ einoModelMode) (model.AgenticModel, error) { + built = append(built, oa.Model) + return &fakeAgenticGateModel{}, nil + }, + nil, + "agentic", + nil, + "eino_single_agentic", + "conv-1", + ) + if err != nil { + t.Fatalf("newEinoAgenticModelFailoverConfig: %v", err) + } + if cfg == nil || cfg.MaxRetries != 1 { + t.Fatalf("agentic failover cfg = %#v, want max retries 1", cfg) + } + m, msgs, err := cfg.GetFailoverModel(context.Background(), &adk.FailoverContext[*schema.AgenticMessage]{FailoverAttempt: 1}) + if err != nil || m == nil || msgs != nil { + t.Fatalf("GetFailoverModel = (%v, %v, %v)", m, msgs, err) + } + if len(built) != 1 || built[0] != "fallback-1" { + t.Fatalf("built models = %v, want [fallback-1]", built) + } + if !cfg.ShouldFailover(context.Background(), nil, &adk.RetryExhaustedError{LastErr: errors.New("upstream returned 503"), TotalRetries: 4}) { + t.Fatal("retry-exhausted transient agentic error should fail over") + } + if cfg.ShouldFailover(context.Background(), nil, &adk.RetryExhaustedError{LastErr: errors.New("invalid api key"), TotalRetries: 4}) { + t.Fatal("retry-exhausted permanent agentic error should not fail over") + } +} + +func TestNewEinoAgenticModelFailoverConfigEmitsProgressEvent(t *testing.T) { + t.Parallel() + appCfg := &config.Config{ + OpenAI: config.OpenAIConfig{Provider: "openai", APIKey: "k1", BaseURL: "https://api.example/v1", Model: "primary"}, + AI: config.AIConfig{Channels: map[string]config.AIChannelConfig{ + "fb1": {Provider: "openai", APIKey: "k2", BaseURL: "https://api.example/v1", Model: "fallback-1"}, + }}, + } + var events []struct { + eventType string + message string + data interface{} + } + cfg, err := newEinoAgenticModelFailoverConfig( + context.Background(), + appCfg, + &config.MultiAgentEinoMiddlewareConfig{ModelFailoverChannels: []string{"fb1"}}, + einoModelModeNormal, + func(_ context.Context, _ config.OpenAIConfig, _ einoModelMode) (model.AgenticModel, error) { + return &fakeAgenticGateModel{}, nil + }, + nil, + "agentic", + func(eventType, message string, data interface{}) { + events = append(events, struct { + eventType string + message string + data interface{} + }{eventType: eventType, message: message, data: data}) + }, + "eino_single_agentic", + "conv-1", + ) + if err != nil { + t.Fatalf("newEinoAgenticModelFailoverConfig: %v", err) + } + if _, _, err := cfg.GetFailoverModel(context.Background(), &adk.FailoverContext[*schema.AgenticMessage]{FailoverAttempt: 1}); err != nil { + t.Fatalf("GetFailoverModel: %v", err) + } + if len(events) != 1 || events[0].eventType != "eino_model_failover" { + t.Fatalf("events = %#v, want one eino_model_failover", events) + } + payload, ok := events[0].data.(map[string]interface{}) + if !ok { + t.Fatalf("event payload type = %T", events[0].data) + } + if payload["conversationId"] != "conv-1" || payload["orchestration"] != "eino_single_agentic" || payload["channel"] != "fb1" || payload["model"] != "fallback-1" { + t.Fatalf("payload = %#v", payload) + } +} + +func TestNewEinoAgenticChatModelFactoryBuildsOpenAIBackend(t *testing.T) { + t.Parallel() + factory := newEinoAgenticChatModelFactory(newEinoBaseHTTPClient(), nil, nil) + m, err := factory(context.Background(), config.OpenAIConfig{ + Provider: "openai", + APIKey: "test-key", + BaseURL: "https://api.example/v1", + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + }, + }, einoModelModeNormal) + if err != nil { + t.Fatalf("agentic factory: %v", err) + } + if m == nil { + t.Fatal("agentic factory returned nil model") + } + gate := evaluateEinoAgenticModelGate(agenticModelGateFactory(factory, config.OpenAIConfig{ + Provider: "openai", + APIKey: "test-key", + BaseURL: "https://api.example/v1", + Model: "gpt-4o-mini", + }, einoModelModeNormal), einoAgenticRuntimeSupportV0914()) + if !gate.Ready { + t.Fatalf("gate = %#v, want ready with buildable agentic backend", gate) + } +} + +func TestNewEinoAgenticChatModelFactoryBuildsNativeClaudeBackend(t *testing.T) { + t.Parallel() + factory := newEinoAgenticChatModelFactory(newEinoBaseHTTPClient(), nil, nil) + m, err := factory(context.Background(), config.OpenAIConfig{ + Provider: "claude", + APIKey: "test-key", + BaseURL: "https://api.anthropic.com/v1", + Model: "claude-sonnet-4", + }, einoModelModeNormal) + if err != nil { + t.Fatalf("claude agentic factory: %v", err) + } + 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) + } + gate := evaluateEinoAgenticModelGate(agenticModelGateFactory(factory, config.OpenAIConfig{ + Provider: "claude", + APIKey: "test-key", + BaseURL: "https://api.anthropic.com/v1", + Model: "claude-sonnet-4", + }, einoModelModeNormal), einoAgenticRuntimeSupportV0914()) + if !gate.Ready { + t.Fatalf("gate = %#v, want ready with native Claude backend", gate) + } +} + +func TestNewEinoToolCallingChatModelFactoryUsesNativeClaudeAdapter(t *testing.T) { + t.Parallel() + factory := newEinoToolCallingChatModelFactory(newEinoBaseHTTPClient(), nil, nil) + m, err := factory(context.Background(), config.OpenAIConfig{ + Provider: "claude", + APIKey: "test-key", + BaseURL: "https://api.anthropic.com", + Model: "claude-sonnet-4", + }, einoModelModePlanner) + if err != nil { + t.Fatalf("claude planner factory: %v", err) + } + if _, ok := m.(*agenticToolCallingChatModelAdapter); !ok { + t.Fatalf("claude planner factory returned %T, want native agentic adapter", m) + } +} + +func TestEinoNativeRetryErrorsDoNotTriggerRunLevelTransientRetry(t *testing.T) { + t.Parallel() + err := &adk.WillRetryError{ErrStr: "HTTP 429 Too Many Requests", RetryAttempt: 1} + if isEinoTransientRunError(err) { + t.Fatal("WillRetryError should be observed, not treated as a run-level transient failure") + } + exhausted := &adk.RetryExhaustedError{LastErr: errors.New("HTTP 429 Too Many Requests"), TotalRetries: 4} + if isEinoTransientRunError(exhausted) { + t.Fatal("RetryExhaustedError should not trigger a second run-level retry layer") + } + if got := unwrapEinoRetryExhausted(exhausted); got == exhausted { + t.Fatal("unwrapEinoRetryExhausted should return the underlying model error") + } +} diff --git a/internal/multiagent/eino_model_rewrite_pipeline.go b/internal/multiagent/eino_model_rewrite_pipeline.go new file mode 100644 index 00000000..aabd3c1d --- /dev/null +++ b/internal/multiagent/eino_model_rewrite_pipeline.go @@ -0,0 +1,38 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" +) + +func applyBeforeModelRewriteHandlers( + ctx context.Context, + msgs []adk.Message, + handlers []adk.ChatModelAgentMiddleware, +) ([]adk.Message, error) { + if len(msgs) == 0 || len(handlers) == 0 { + return msgs, nil + } + state := &adk.ChatModelAgentState{Messages: msgs} + modelCtx := &adk.ModelContext{} + curCtx := ctx + for _, h := range handlers { + if h == nil { + continue + } + nextCtx, nextState, err := h.BeforeModelRewriteState(curCtx, state, modelCtx) + if err != nil { + return nil, fmt.Errorf("before model rewrite: %w", err) + } + if nextCtx != nil { + curCtx = nextCtx + } + if nextState != nil { + state = nextState + } + } + return state.Messages, nil +} + diff --git a/internal/multiagent/eino_native_cancel.go b/internal/multiagent/eino_native_cancel.go new file mode 100644 index 00000000..605390b6 --- /dev/null +++ b/internal/multiagent/eino_native_cancel.go @@ -0,0 +1,102 @@ +package multiagent + +import ( + "context" + "errors" + "time" + + "github.com/cloudwego/eino/adk" +) + +const ( + einoNativeCancelImmediateWait = 1200 * time.Millisecond + einoNativeCancelSafePointWait = 3500 * time.Millisecond + einoNativeCancelSafePointTTL = 3 * time.Second +) + +type agentRuntimeCancelRegistrarKey struct{} +type agentTurnLoopInterruptRegistrarKey struct{} + +// AgentRuntimeCancelRegistrar binds the currently active Eino ADK cancel hook +// into the host task manager. The hook returns true when Eino accepted and +// handled the cancel request, so the host can avoid canceling the parent context. +type AgentRuntimeCancelRegistrar func(cancel func(error) bool) (unregister func()) + +// WithAgentRuntimeCancelRegistrar lets the HTTP/task layer trigger Eino's native +// Agent Cancel before falling back to the existing context cancellation path. +func WithAgentRuntimeCancelRegistrar(ctx context.Context, registrar AgentRuntimeCancelRegistrar) context.Context { + if ctx == nil || registrar == nil { + return ctx + } + return context.WithValue(ctx, agentRuntimeCancelRegistrarKey{}, registrar) +} + +func agentRuntimeCancelRegistrarFromContext(ctx context.Context) AgentRuntimeCancelRegistrar { + if ctx == nil { + return nil + } + if v, ok := ctx.Value(agentRuntimeCancelRegistrarKey{}).(AgentRuntimeCancelRegistrar); ok { + return v + } + return nil +} + +// AgentTurnLoopInterruptRegistrar binds a conversation-level TurnLoop interrupt +// pusher into the host task manager. The pusher receives the user supplied note +// and returns true when the note was accepted by the loop. +type AgentTurnLoopInterruptRegistrar func(push func(note string) bool) (unregister func()) + +// WithAgentTurnLoopInterruptRegistrar lets the HTTP/task layer enqueue a user +// supplement into an active Eino TurnLoop before falling back to cancellation. +func WithAgentTurnLoopInterruptRegistrar(ctx context.Context, registrar AgentTurnLoopInterruptRegistrar) context.Context { + if ctx == nil || registrar == nil { + return ctx + } + return context.WithValue(ctx, agentTurnLoopInterruptRegistrarKey{}, registrar) +} + +func agentTurnLoopInterruptRegistrarFromContext(ctx context.Context) AgentTurnLoopInterruptRegistrar { + if ctx == nil { + return nil + } + if v, ok := ctx.Value(agentTurnLoopInterruptRegistrarKey{}).(AgentTurnLoopInterruptRegistrar); ok { + return v + } + return nil +} + +func requestEinoNativeAgentCancel(cancelFn adk.AgentCancelFunc, cause error) (waitErr error, submitted bool, handled bool) { + if cancelFn == nil { + return nil, false, false + } + opts, waitFor := einoNativeCancelOptions(cause) + handle, submitted := cancelFn(opts...) + if !submitted || handle == nil { + return nil, submitted, false + } + waitCh := make(chan error, 1) + go func() { + waitCh <- handle.Wait() + }() + select { + case err := <-waitCh: + handled := err == nil || errors.Is(err, adk.ErrCancelTimeout) || errors.Is(err, adk.ErrExecutionEnded) + return err, submitted, handled + case <-time.After(waitFor): + return context.DeadlineExceeded, submitted, false + } +} + +func einoNativeCancelOptions(cause error) ([]adk.AgentCancelOption, time.Duration) { + if errors.Is(cause, ErrInterruptContinue) { + return []adk.AgentCancelOption{ + adk.WithAgentCancelMode(adk.CancelAfterChatModel | adk.CancelAfterToolCalls), + adk.WithAgentCancelTimeout(einoNativeCancelSafePointTTL), + adk.WithRecursive(), + }, einoNativeCancelSafePointWait + } + return []adk.AgentCancelOption{ + adk.WithAgentCancelMode(adk.CancelImmediate), + adk.WithRecursive(), + }, einoNativeCancelImmediateWait +} diff --git a/internal/multiagent/eino_native_cancel_test.go b/internal/multiagent/eino_native_cancel_test.go new file mode 100644 index 00000000..e3b4c5a7 --- /dev/null +++ b/internal/multiagent/eino_native_cancel_test.go @@ -0,0 +1,27 @@ +package multiagent + +import ( + "context" + "testing" +) + +func TestEinoNativeCancelOptionsByCause(t *testing.T) { + fullStopOpts, fullStopWait := einoNativeCancelOptions(context.Canceled) + if len(fullStopOpts) != 2 { + t.Fatalf("full stop options: got %d want 2", len(fullStopOpts)) + } + if fullStopWait != einoNativeCancelImmediateWait { + t.Fatalf("full stop wait: got %v want %v", fullStopWait, einoNativeCancelImmediateWait) + } + + interruptOpts, interruptWait := einoNativeCancelOptions(ErrInterruptContinue) + if len(interruptOpts) != 3 { + t.Fatalf("interrupt options: got %d want 3", len(interruptOpts)) + } + if interruptWait != einoNativeCancelSafePointWait { + t.Fatalf("interrupt wait: got %v want %v", interruptWait, einoNativeCancelSafePointWait) + } + if interruptWait <= einoNativeCancelSafePointTTL { + t.Fatalf("interrupt wait must allow the Eino safe-point timeout to elapse: wait=%v ttl=%v", interruptWait, einoNativeCancelSafePointTTL) + } +} diff --git a/internal/multiagent/eino_native_model_retry_progress.go b/internal/multiagent/eino_native_model_retry_progress.go new file mode 100644 index 00000000..c41890e9 --- /dev/null +++ b/internal/multiagent/eino_native_model_retry_progress.go @@ -0,0 +1,41 @@ +package multiagent + +import ( + "fmt" + + "github.com/cloudwego/eino/adk" + "go.uber.org/zap" +) + +func emitEinoNativeModelRetryProgress( + conversationID, orchMode string, + willRetry *adk.WillRetryError, + progress func(eventType, message string, data interface{}), + logger *zap.Logger, + runErr error, +) bool { + if willRetry == nil { + return false + } + if progress != nil { + reason := "" + if willRetry.RejectReason() != nil { + reason = fmt.Sprint(willRetry.RejectReason()) + } + progress("eino_model_retry", "模型调用遇到临时问题,Eino 正在原生重试…", map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "orchestration": orchMode, + "attempt": willRetry.RetryAttempt, + "reason": reason, + "error": willRetry.Error(), + }) + } + if logger != nil { + logger.Warn("eino native model retry event", + zap.String("orchestration", orchMode), + zap.Int("attempt", willRetry.RetryAttempt), + zap.Error(runErr)) + } + return true +} diff --git a/internal/multiagent/eino_native_model_retry_progress_test.go b/internal/multiagent/eino_native_model_retry_progress_test.go new file mode 100644 index 00000000..c3030b2c --- /dev/null +++ b/internal/multiagent/eino_native_model_retry_progress_test.go @@ -0,0 +1,85 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/adk" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func TestEmitEinoNativeModelRetryProgress(t *testing.T) { + willRetry := &adk.WillRetryError{ + ErrStr: "HTTP 429 Too Many Requests", + RetryAttempt: 2, + } + var gotType, gotMessage string + var gotData map[string]interface{} + called := emitEinoNativeModelRetryProgress("conv-1", "deep_agent", willRetry, func(eventType, message string, data interface{}) { + gotType = eventType + gotMessage = message + var ok bool + gotData, ok = data.(map[string]interface{}) + if !ok { + t.Fatalf("progress data type = %T, want map[string]interface{}", data) + } + }, nil, willRetry) + if !called { + t.Fatal("called = false, want true") + } + if gotType != "eino_model_retry" { + t.Fatalf("event type = %q, want eino_model_retry", gotType) + } + if gotMessage != "模型调用遇到临时问题,Eino 正在原生重试…" { + t.Fatalf("message = %q", gotMessage) + } + assertNativeRetryMapValue(t, gotData, "conversationId", "conv-1") + assertNativeRetryMapValue(t, gotData, "source", "eino") + assertNativeRetryMapValue(t, gotData, "orchestration", "deep_agent") + assertNativeRetryMapValue(t, gotData, "attempt", 2) + assertNativeRetryMapValue(t, gotData, "reason", "") + assertNativeRetryMapValue(t, gotData, "error", "HTTP 429 Too Many Requests") +} + +func TestEmitEinoNativeModelRetryProgressNilSafe(t *testing.T) { + calledProgress := false + called := emitEinoNativeModelRetryProgress("conv-1", "deep_agent", nil, func(string, string, interface{}) { + calledProgress = true + }, nil, nil) + if called { + t.Fatal("called = true, want false") + } + if calledProgress { + t.Fatal("progress called for nil willRetry") + } +} + +func TestEmitEinoNativeModelRetryProgressLogsEvent(t *testing.T) { + core, logs := observer.New(zap.WarnLevel) + logger := zap.New(core) + willRetry := &adk.WillRetryError{ + ErrStr: "HTTP 500", + RetryAttempt: 3, + } + + emitEinoNativeModelRetryProgress("conv-1", "single_agent", willRetry, nil, logger, willRetry) + + entry := logs.FilterMessage("eino native model retry event").TakeAll() + if len(entry) != 1 { + t.Fatalf("log count = %d, want 1", len(entry)) + } + fields := entry[0].ContextMap() + if fields["orchestration"] != "single_agent" { + t.Fatalf("orchestration field = %v", fields["orchestration"]) + } + if fields["attempt"] != int64(3) { + t.Fatalf("attempt field = %v", fields["attempt"]) + } +} + +func assertNativeRetryMapValue(t *testing.T, data map[string]interface{}, key string, want interface{}) { + t.Helper() + if got := data[key]; got != want { + t.Fatalf("%s = %v, want %v", key, got, want) + } +} diff --git a/internal/multiagent/eino_orchestration.go b/internal/multiagent/eino_orchestration.go new file mode 100644 index 00000000..051984e2 --- /dev/null +++ b/internal/multiagent/eino_orchestration.go @@ -0,0 +1,408 @@ +package multiagent + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/prebuilt/planexecute" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// PlanExecuteRootArgs 构建 Eino adk/prebuilt/planexecute 根 Agent 所需参数。 +type PlanExecuteRootArgs struct { + MainToolCallingModel model.ToolCallingChatModel + AgenticExecModel model.AgenticModel + OrchInstruction string + ToolsCfg adk.ToolsConfig + ExecMaxIter int + LoopMaxIter int + // AppCfg / Logger 非空时为 Executor 挂载与 Deep/Supervisor 一致的 Eino summarization 中间件。 + AppCfg *config.Config + MwCfg *config.MultiAgentEinoMiddlewareConfig + // ConversationID is used for transcript/isolation paths in middleware. + ConversationID string + DB *database.DB + ProjectID string + Logger *zap.Logger + // ModelName is used for model input token estimation logs. + ModelName string + // AgenticExecPreMiddlewares 是由 prependEinoAgenticMiddlewares 构建的前置中间件。 + AgenticExecPreMiddlewares []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + AgenticSkillMiddleware adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + AgenticFilesystemMiddleware adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + // PlannerReplannerRewriteHandlers applies BeforeModelRewriteState pipeline for planner/replanner input. + PlannerReplannerRewriteHandlers []adk.ChatModelAgentMiddleware + // ModelFacingTrace 可选:由 Executor Handlers 链末尾写入,供 last_react 与 summarization 后上下文对齐。 + ModelFacingTrace *modelFacingTraceHolder + AgenticModelRetryConfig *adk.TypedModelRetryConfig[*schema.AgenticMessage] + AgenticModelFailoverConfig *adk.ModelFailoverConfig[*schema.AgenticMessage] +} + +// NewPlanExecuteRoot 返回 plan → execute → replan 预置编排根节点(与 Deep / Supervisor 并列)。 +func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.ResumableAgent, error) { + if a == nil { + return nil, fmt.Errorf("plan_execute: args 为空") + } + if a.MainToolCallingModel == nil || a.AgenticExecModel == nil { + return nil, fmt.Errorf("plan_execute: 模型为空") + } + tcm, ok := interface{}(a.MainToolCallingModel).(model.ToolCallingChatModel) + if !ok { + return nil, fmt.Errorf("plan_execute: 主模型需实现 ToolCallingChatModel") + } + plannerCfg := &planexecute.PlannerConfig{ + ToolCallingChatModel: tcm, + NewPlan: newLenientPlan, + } + if fn := planExecutePlannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers); fn != nil { + plannerCfg.GenInputFn = fn + } + planner, err := planexecute.NewPlanner(ctx, plannerCfg) + if err != nil { + return nil, fmt.Errorf("plan_execute planner: %w", err) + } + replanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{ + ChatModel: tcm, + GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers), + NewPlan: newLenientPlan, + }) + if err != nil { + return nil, fmt.Errorf("plan_execute replanner: %w", err) + } + + var executor adk.Agent + agenticExecHandlers, herr := buildPlanExecuteAgenticExecutorHandlers(ctx, a) + if herr != nil { + return nil, herr + } + executor, err = newPlanExecuteAgenticExecutor(ctx, &planexecute.ExecutorConfig{ + ToolsConfig: a.ToolsCfg, + MaxIterations: a.ExecMaxIter, + GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID), + }, a.AgenticExecModel, agenticExecHandlers, a.AgenticModelRetryConfig, a.AgenticModelFailoverConfig) + if err != nil { + return nil, fmt.Errorf("plan_execute executor: %w", err) + } + loopMax := a.LoopMaxIter + if loopMax <= 0 { + loopMax = 10 + } + return planexecute.New(ctx, &planexecute.Config{ + Planner: planner, + Executor: executor, + Replanner: replanner, + MaxIterations: loopMax, + }) +} + +func buildPlanExecuteAgenticExecutorHandlers(ctx context.Context, a *PlanExecuteRootArgs) ([]adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage], error) { + if a == nil { + return nil, fmt.Errorf("plan_execute: args 为空") + } + var execHandlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + if len(a.AgenticExecPreMiddlewares) > 0 { + execHandlers = append(execHandlers, a.AgenticExecPreMiddlewares...) + } + if a.AgenticFilesystemMiddleware != nil { + execHandlers = append(execHandlers, a.AgenticFilesystemMiddleware) + } + if a.AgenticSkillMiddleware != nil { + execHandlers = append(execHandlers, a.AgenticSkillMiddleware) + } + if a.AppCfg != nil { + sumMw, sumErr := newEinoAgenticSummarizationMiddleware(ctx, a.AgenticExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger) + if sumErr != nil { + return nil, fmt.Errorf("plan_execute agentic executor summarization: %w", sumErr) + } + execHandlers = appendEinoAgenticChatModelTailMiddlewares(execHandlers, einoChatModelTailConfig{ + logger: a.Logger, + phase: "plan_execute_executor", + agenticSummarization: sumMw, + modelName: a.ModelName, + maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens, + toolMaxBytes: toolMaxBytesFromMW(a.MwCfg), + conversationID: a.ConversationID, + trace: a.ModelFacingTrace, + middlewareConfig: a.MwCfg, + }) + } + return execHandlers, nil +} + +// planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。 +// 返回 nil 时 Eino 使用内置默认 planner prompt。 +func planExecutePlannerGenInput( + orchInstruction string, + appCfg *config.Config, + mwCfg *config.MultiAgentEinoMiddlewareConfig, + logger *zap.Logger, + modelName string, + conversationID string, + rewriteHandlers []adk.ChatModelAgentMiddleware, +) planexecute.GenPlannerModelInputFn { + oi := strings.TrimSpace(orchInstruction) + if oi == "" && appCfg == nil { + return nil + } + return func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) { + userInput = capPlanExecuteUserInputMessages(userInput, appCfg, mwCfg) + msgs := make([]adk.Message, 0, len(userInput)) + msgs = append(msgs, userInput...) + if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 { + msgs = rewritten + } + msgs = normalizeSingleLeadingSystemMessage(msgs, oi) + logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_planner", msgs) + return msgs, nil + } +} + +func planExecuteExecutorGenInput( + orchInstruction string, + appCfg *config.Config, + mwCfg *config.MultiAgentEinoMiddlewareConfig, + logger *zap.Logger, + modelName string, + conversationID string, +) planexecute.GenModelInputFn { + oi := strings.TrimSpace(orchInstruction) + return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) { + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + userMsgs, err := planexecute.ExecutorPrompt.Format(ctx, map[string]any{ + "input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)), + "plan": string(planContent), + "executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg), + "step": in.Plan.FirstStep(), + }) + if err != nil { + return nil, err + } + userMsgs = normalizeSingleLeadingSystemMessage(userMsgs, oi) + logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_executor_gen_input", userMsgs) + return userMsgs, nil + } +} + +func planExecuteFormatInput(input []adk.Message) string { + var sb strings.Builder + for _, msg := range input { + sb.WriteString(msg.Content) + sb.WriteString("\n") + } + return sb.String() +} + +func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string { + capped := capPlanExecuteExecutedStepsWithConfig(results, mwCfg) + return renderPlanExecuteStepsByBudget(capped, appCfg, mwCfg) +} + +// planExecuteReplannerGenInput 与 Eino 默认 Replanner 输入一致,但 executed_steps 经 cap 后再写入 prompt, +// 且在 orchInstruction 非空时 prepend SystemMessage 使 replanner 也能接收全局指令。 +func planExecuteReplannerGenInput( + orchInstruction string, + appCfg *config.Config, + mwCfg *config.MultiAgentEinoMiddlewareConfig, + logger *zap.Logger, + modelName string, + conversationID string, + rewriteHandlers []adk.ChatModelAgentMiddleware, +) planexecute.GenModelInputFn { + oi := strings.TrimSpace(orchInstruction) + return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) { + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + msgs, err := planexecute.ReplannerPrompt.Format(ctx, map[string]any{ + "plan": string(planContent), + "input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)), + "executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg), + "plan_tool": planexecute.PlanToolInfo.Name, + "respond_tool": planexecute.RespondToolInfo.Name, + }) + if err != nil { + return nil, err + } + if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 { + msgs = rewritten + } + msgs = normalizeSingleLeadingSystemMessage(msgs, oi) + logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_replanner", msgs) + return msgs, nil + } +} + +// normalizeSingleLeadingSystemMessage enforces a provider-friendly message shape: +// exactly one system message at index 0 (when any system context exists). +// For strict OpenAI-compatible backends (e.g. qwen/vllm templates), this avoids +// "System message must be at the beginning" caused by multiple/disordered system messages. +func normalizeSingleLeadingSystemMessage(msgs []adk.Message, extraSystem string) []adk.Message { + extraSystem = strings.TrimSpace(extraSystem) + if len(msgs) == 0 { + if extraSystem == "" { + return msgs + } + return []adk.Message{schema.SystemMessage(extraSystem)} + } + + systemParts := make([]string, 0, 2) + if extraSystem != "" { + systemParts = append(systemParts, extraSystem) + } + nonSystem := make([]adk.Message, 0, len(msgs)) + for _, msg := range msgs { + if msg == nil { + continue + } + if msg.Role == schema.System { + if s := strings.TrimSpace(msg.Content); s != "" { + systemParts = append(systemParts, s) + } + continue + } + nonSystem = append(nonSystem, msg) + } + if len(systemParts) == 0 { + return nonSystem + } + out := make([]adk.Message, 0, len(nonSystem)+1) + out = append(out, schema.SystemMessage(strings.Join(systemParts, "\n\n"))) + out = append(out, nonSystem...) + return out +} + +func capPlanExecuteUserInputMessages(input []adk.Message, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message { + if len(input) == 0 { + return input + } + maxTotal := 120000 + modelName := "gpt-4o" + if appCfg != nil { + if appCfg.OpenAI.MaxTotalTokens > 0 { + maxTotal = appCfg.OpenAI.MaxTotalTokens + } + if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" { + modelName = m + } + } + // Reserve most tokens for planner/replanner prompt and tool schema. + ratio := 0.35 + if mwCfg != nil { + ratio = mwCfg.PlanExecuteUserInputBudgetRatioEffective() + } + budget := int(float64(maxTotal) * ratio) + if budget < 4096 { + budget = 4096 + } + tc := agent.NewTikTokenCounter() + out := make([]adk.Message, 0, len(input)) + used := 0 + for i := len(input) - 1; i >= 0; i-- { + msg := input[i] + if msg == nil { + continue + } + n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content) + if err != nil { + n = (len(msg.Content) + 3) / 4 + } + if n <= 0 { + n = 1 + } + if used+n > budget { + break + } + used += n + out = append(out, msg) + } + for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 { + out[i], out[j] = out[j], out[i] + } + if len(out) == 0 { + // Keep the latest user message at least. + return []adk.Message{input[len(input)-1]} + } + return out +} + +func renderPlanExecuteStepsByBudget(steps []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string { + if len(steps) == 0 { + return "" + } + maxTotal := 120000 + modelName := "gpt-4o" + if appCfg != nil { + if appCfg.OpenAI.MaxTotalTokens > 0 { + maxTotal = appCfg.OpenAI.MaxTotalTokens + } + if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" { + modelName = m + } + } + ratio := 0.2 + if mwCfg != nil { + ratio = mwCfg.PlanExecuteExecutedStepsBudgetRatioEffective() + } + budget := int(float64(maxTotal) * ratio) + if budget < 3072 { + budget = 3072 + } + tc := agent.NewTikTokenCounter() + var kept []string + used := 0 + skipped := 0 + for i := len(steps) - 1; i >= 0; i-- { + block := fmt.Sprintf("Step: %s\nResult: %s\n\n", steps[i].Step, steps[i].Result) + n, err := tc.Count(modelName, block) + if err != nil { + n = (len(block) + 3) / 4 + } + if n <= 0 { + n = 1 + } + if used+n > budget { + skipped = i + 1 + break + } + used += n + kept = append(kept, block) + } + var sb strings.Builder + if skipped > 0 { + sb.WriteString(fmt.Sprintf("Earlier executed steps omitted due to context budget: %d steps.\n\n", skipped)) + } + for i := len(kept) - 1; i >= 0; i-- { + sb.WriteString(kept[i]) + } + return sb.String() +} + +// planExecuteStreamsMainAssistant 将规划/执行/重规划各阶段助手流式输出映射到主对话区。 +func planExecuteStreamsMainAssistant(agent string) bool { + if agent == "" { + return true + } + switch agent { + case "planner", "executor", "replanner", "execute_replan", "plan_execute_replan": + return true + default: + return false + } +} + +func planExecuteEinoRoleTag(agent string) string { + _ = agent + return "orchestrator" +} diff --git a/internal/multiagent/eino_orchestration_system_message_test.go b/internal/multiagent/eino_orchestration_system_message_test.go new file mode 100644 index 00000000..2cb32cfc --- /dev/null +++ b/internal/multiagent/eino_orchestration_system_message_test.go @@ -0,0 +1,45 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestNormalizeSingleLeadingSystemMessage_MergesMultipleSystems(t *testing.T) { + in := []adk.Message{ + schema.SystemMessage("sys-1"), + schema.UserMessage("u1"), + schema.SystemMessage("sys-2"), + schema.AssistantMessage("a1", nil), + } + out := normalizeSingleLeadingSystemMessage(in, "orch") + if len(out) != 3 { + t.Fatalf("unexpected output length: got %d want 3", len(out)) + } + if out[0].Role != schema.System { + t.Fatalf("first message role must be system, got %s", out[0].Role) + } + if got := out[0].Content; got != "orch\n\nsys-1\n\nsys-2" { + t.Fatalf("unexpected merged system content: %q", got) + } + if out[1].Role != schema.User || out[2].Role != schema.Assistant { + t.Fatalf("non-system message order changed unexpectedly") + } +} + +func TestNormalizeSingleLeadingSystemMessage_NoSystemKeepsFlow(t *testing.T) { + in := []adk.Message{ + schema.UserMessage("u1"), + schema.AssistantMessage("a1", nil), + } + out := normalizeSingleLeadingSystemMessage(in, "") + if len(out) != 2 { + t.Fatalf("unexpected output length: got %d want 2", len(out)) + } + if out[0].Role != schema.User || out[1].Role != schema.Assistant { + t.Fatalf("message order changed unexpectedly") + } +} + diff --git a/internal/multiagent/eino_pending_tool_calls_test.go b/internal/multiagent/eino_pending_tool_calls_test.go new file mode 100644 index 00000000..c241cf11 --- /dev/null +++ b/internal/multiagent/eino_pending_tool_calls_test.go @@ -0,0 +1,81 @@ +package multiagent + +import ( + "errors" + "testing" +) + +func TestEinoPendingToolCallsPopNextForAgentSkipsRemovedIDs(t *testing.T) { + p := newEinoPendingToolCalls("conv", nil) + p.Mark(toolCallPendingInfo{ToolCallID: "call-1", ToolName: "first", EinoAgent: "agent"}) + p.Mark(toolCallPendingInfo{ToolCallID: "call-2", ToolName: "second", EinoAgent: "agent"}) + p.RemoveByID("call-1") + + got, ok := p.PopNextForAgent("agent") + if !ok { + t.Fatal("expected pending tool call") + } + if got.ToolCallID != "call-2" { + t.Fatalf("toolCallID = %q, want call-2", got.ToolCallID) + } + if p.Count() != 0 { + t.Fatalf("pending count = %d, want 0", p.Count()) + } +} + +func TestEinoPendingToolCallsPopAny(t *testing.T) { + p := newEinoPendingToolCalls("conv", nil) + p.Mark(toolCallPendingInfo{ToolCallID: "call-1", ToolName: "tool"}) + + got, ok := p.PopAny() + if !ok || got.ToolCallID != "call-1" { + t.Fatalf("PopAny = %#v ok=%v", got, ok) + } + if _, ok := p.PopAny(); ok { + t.Fatal("PopAny should be empty after first pop") + } +} + +func TestEinoPendingToolCallsFlushAsFailedEmitsAndClears(t *testing.T) { + var events []struct { + eventType string + message string + data map[string]interface{} + } + p := newEinoPendingToolCalls("conv-1", func(eventType, message string, data interface{}) { + m, _ := data.(map[string]interface{}) + events = append(events, struct { + eventType string + message string + data map[string]interface{} + }{eventType: eventType, message: message, data: m}) + }) + p.Mark(toolCallPendingInfo{ + ToolCallID: "call-err", + ToolName: "", + EinoAgent: "agent", + EinoRole: "sub", + }) + + p.FlushAsFailed(errors.New("boom")) + + if p.Count() != 0 { + t.Fatalf("pending count = %d, want 0", p.Count()) + } + if len(events) != 1 { + t.Fatalf("events = %#v, want one", events) + } + ev := events[0] + if ev.eventType != "tool_result" || ev.message != "工具结果 (unknown)" { + t.Fatalf("event = %#v", ev) + } + if ev.data["toolCallId"] != "call-err" || + ev.data["conversationId"] != "conv-1" || + ev.data["einoAgent"] != "agent" || + ev.data["einoRole"] != "sub" || + ev.data["isError"] != true || + ev.data["success"] != false || + ev.data["result"] != "boom" { + t.Fatalf("payload = %#v", ev.data) + } +}