diff --git a/internal/multiagent/runner.go b/internal/multiagent/runner.go new file mode 100644 index 00000000..d62bf3ca --- /dev/null +++ b/internal/multiagent/runner.go @@ -0,0 +1,1102 @@ +// Package multiagent 使用 CloudWeGo Eino adk/prebuilt(deep / plan_execute / supervisor)编排多代理,MCP 工具经 einomcp 桥接到现有 Agent。 +package multiagent + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "sync" + "sync/atomic" + "unicode/utf8" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/agents" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/einomcp" + "cyberstrike-ai/internal/project" + "cyberstrike-ai/internal/reasoning" + "cyberstrike-ai/internal/security" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/prebuilt/deep" + "github.com/cloudwego/eino/adk/prebuilt/supervisor" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// RunResult 与单 Agent 循环结果字段对齐,便于复用存储与 SSE 收尾逻辑。 +type RunResult struct { + Response string + MCPExecutionIDs []string + LastAgentTraceInput string // 已序列化的消息带(JSON):原生循环或 Eino 均写入,供续跑/攻击链等恢复上下文 + LastAgentTraceOutput string // 本轮助手侧对外展示文本(摘要或最终回复) + Finalized bool + Status string + CompletionReason string + EvidenceVerified bool + EvidenceRefs []string + PendingExecutionIDs []string + MissingChecks []string +} + +// toolCallPendingInfo tracks a tool_call emitted to the UI so we can later +// correlate tool_result events (even when the framework omits ToolCallID) and +// avoid leaving the UI stuck in "running" state on recoverable errors. +type toolCallPendingInfo struct { + ToolCallID string + ToolName string + Arguments map[string]interface{} + EinoAgent string + EinoRole string +} + +var fallbackToolCallSequence atomic.Uint64 + +// RunDeepAgent 使用 Eino 多代理预置编排执行一轮对话(deep / plan_execute / supervisor;流式事件通过 progress 回调输出)。 +// orchestrationOverride 非空时优先(如聊天/WebShell 请求体);否则用 multi_agent.orchestration(遗留 yaml);皆空则按 deep。 +// reasoningClient 来自 ChatRequest.reasoning;可为 nil(机器人/批量等走全局 openai.reasoning)。 +func RunDeepAgent( + ctx context.Context, + appCfg *config.Config, + ma *config.MultiAgentConfig, + ag *agent.Agent, + db *database.DB, + logger *zap.Logger, + conversationID string, + projectID string, + userMessage string, + history []agent.ChatMessage, + roleTools []string, + progress func(eventType, message string, data interface{}), + agentsMarkdownDir string, + orchestrationOverride string, + reasoningClient *reasoning.ClientIntent, + systemPromptExtra string, +) (*RunResult, error) { + if appCfg == nil || ma == nil || ag == nil { + return nil, fmt.Errorf("multiagent: 配置或 Agent 为空") + } + + runtimeUserMessage := prepareLatestUserMessageForModel(userMessage, appCfg, &ma.EinoMiddleware, conversationID, logger) + + effectiveSubs := ma.SubAgents + var markdownLoad *agents.MarkdownDirLoad + var orch *agents.OrchestratorMarkdown + if strings.TrimSpace(agentsMarkdownDir) != "" { + load, merr := agents.LoadMarkdownAgentsDir(agentsMarkdownDir) + if merr != nil { + if logger != nil { + logger.Warn("加载 agents 目录 Markdown 失败,沿用 config 中的 sub_agents", zap.Error(merr)) + } + } else { + markdownLoad = load + effectiveSubs = agents.MergeYAMLAndMarkdown(ma.SubAgents, load.SubAgents) + orch = load.Orchestrator + } + } + orchMode := config.NormalizeMultiAgentOrchestration(ma.Orchestration) + if o := strings.TrimSpace(orchestrationOverride); o != "" { + orchMode = config.NormalizeMultiAgentOrchestration(o) + } + if orchMode != "plan_execute" && ma.WithoutGeneralSubAgent && len(effectiveSubs) == 0 { + return nil, fmt.Errorf("multi_agent.without_general_sub_agent 为 true 时,必须在 multi_agent.sub_agents 或 agents 目录 Markdown 中配置至少一个子代理") + } + if orchMode == "supervisor" && len(effectiveSubs) == 0 { + return nil, fmt.Errorf("multi_agent.orchestration=supervisor 时需至少配置一个子代理(sub_agents 或 agents 目录 Markdown)") + } + if orchMode == "supervisor" && len(effectiveSubs) == 1 && progress != nil { + progress("progress", "Supervisor 是专家路由模式;当前仅 1 个子代理,专家路由空间有限,仍会继续执行。", map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "orchestration": orchMode, + "kind": "supervisor_boundary_hint", + }) + } + + agenticLoc, agenticSkillMW, agenticFSTools, agenticSkillsRoot, einoErr := prepareEinoAgenticSkills(ctx, appCfg.SkillsDir, ma, logger) + if einoErr != nil { + return nil, einoErr + } + + holder := &einomcp.ConversationHolder{} + holder.Set(conversationID) + + var mcpIDsMu sync.Mutex + var mcpIDs []string + mcpExecBinder := NewMCPExecutionBinder() + recorder := func(id, toolCallID string) { + if id == "" { + return + } + mcpExecBinder.Bind(toolCallID, id) + mcpIDsMu.Lock() + mcpIDs = append(mcpIDs, id) + mcpIDsMu.Unlock() + } + einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder) + + // 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。 + snapshotMCPIDs := func() []string { + mcpIDsMu.Lock() + defer mcpIDsMu.Unlock() + out := make([]string, len(mcpIDs)) + copy(out, mcpIDs) + return out + } + + toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder() + mainDefs := ag.ToolsForRole(roleTools) + + baseHTTPClient := newEinoBaseHTTPClient() + modelFactory := newEinoToolCallingChatModelFactory(baseHTTPClient, reasoningClient, logger) + agenticModelFactory := newEinoAgenticChatModelFactory(baseHTTPClient, reasoningClient, logger) + agenticModelRetryCfg := newEinoAgenticModelRetryConfig(&ma.EinoMiddleware, logger, "multiagent") + agenticModelFailoverCfg, err := newEinoAgenticModelFailoverConfig(ctx, appCfg, &ma.EinoMiddleware, einoModelModeNormal, agenticModelFactory, logger, "multiagent", progress, orchMode, conversationID) + if err != nil { + return nil, err + } + logEinoAgenticModelGate( + logger, + "multiagent", + orchMode, + evaluateEinoAgenticModelGate(agenticModelGateFactory(agenticModelFactory, appCfg.OpenAI, einoModelModeNormal), einoAgenticRuntimeSupportV0914()), + ) + + deepMaxIter := agentMaxIterations(appCfg) + + var subAgents []adk.TypedAgent[*schema.AgenticMessage] + var supervisorSubAgents []adk.Agent + if orchMode != "plan_execute" { + subAgents = make([]adk.TypedAgent[*schema.AgenticMessage], 0, len(effectiveSubs)) + supervisorSubAgents = make([]adk.Agent, 0, len(effectiveSubs)) + for _, sub := range effectiveSubs { + id := strings.TrimSpace(sub.ID) + if id == "" { + return nil, fmt.Errorf("multi_agent.sub_agents 中存在空的 id") + } + name := strings.TrimSpace(sub.Name) + if name == "" { + name = id + } + desc := strings.TrimSpace(sub.Description) + if desc == "" { + desc = fmt.Sprintf("Specialist agent %s for penetration testing workflow.", id) + } + instr := strings.TrimSpace(sub.Instruction) + if instr == "" { + instr = "你是 CyberStrikeAI 中的专业子代理,在授权渗透测试场景下协助完成用户委托的子任务。优先使用可用工具获取证据,回答简洁专业。" + } + + roleTools := sub.RoleTools + bind := strings.TrimSpace(sub.BindRole) + if bind != "" && appCfg.Roles != nil { + if r, ok := appCfg.Roles[bind]; ok && r.Enabled { + if len(roleTools) == 0 && len(r.Tools) > 0 { + roleTools = r.Tools + } + } + } + + subModel, err := agenticModelFactory(ctx, appCfg.OpenAI, einoModelModeNormal) + if err != nil { + return nil, fmt.Errorf("子代理 %q AgenticModel: %w", id, err) + } + + subDefs := ag.ToolsForRole(roleTools) + subTools, err := einomcp.ToolsFromDefinitions(ag, holder, subDefs, recorder, nil, toolInvokeNotify, id) + if err != nil { + return nil, fmt.Errorf("子代理 %q 工具: %w", id, err) + } + + subToolsForCfg, subPre, subToolSearchActive, err := prependEinoAgenticMiddlewares(ctx, &ma.EinoMiddleware, einoMWSub, subTools, agenticLoc, agenticSkillsRoot, conversationID, projectID, logger) + if err != nil { + return nil, fmt.Errorf("子代理 %q eino 中间件: %w", id, err) + } + + subMax := resolveMaxIterations(appCfg, sub.MaxIterations) + + subSumMw, err := newEinoAgenticSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) + if err != nil { + return nil, fmt.Errorf("子代理 %q agentic summarization 中间件: %w", id, err) + } + + var subHandlers []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + if len(subPre) > 0 { + subHandlers = append(subHandlers, subPre...) + } + if agenticSkillMW != nil { + if agenticFSTools && agenticLoc != nil { + subFs, fsErr := subAgentAgenticFilesystemMiddleware(ctx, agenticLoc, toolInvokeNotify, id, conversationID, projectID, ma.EinoMiddleware.ReductionRootDir, toolMaxBytesFromMW(&ma.EinoMiddleware), mcpExecBinder, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil) + if fsErr != nil { + return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr) + } + subHandlers = append(subHandlers, subFs) + } + subHandlers = append(subHandlers, agenticSkillMW) + } + subHandlers = appendEinoAgenticChatModelTailMiddlewares(subHandlers, einoChatModelTailConfig{ + logger: logger, + phase: "sub_agent:" + id, + agenticSummarization: subSumMw, + modelName: appCfg.OpenAI.Model, + maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, + toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware), + conversationID: conversationID, + middlewareConfig: &ma.EinoMiddleware, + }) + + subInstrFinal := project.AppendVisionImageAnalysisIfReady(instr, appCfg.Vision.Ready()) + subInstrFinal = injectToolNamesOnlyInstruction(ctx, subInstrFinal, subTools, subToolSearchActive) + if logger != nil { + subNames := collectToolNames(ctx, subTools) + mountedNames := collectToolNames(ctx, subToolsForCfg) + logger.Info("eino tool-name injection", + zap.String("scope", "sub_agent"), + zap.String("agent", id), + zap.Int("tool_names", len(subNames)), + zap.Int("mounted_tool_names", len(mountedNames)), + zap.Bool("tool_search_middleware", subToolSearchActive), + ) + } + sa, err := newEinoAgenticChatModelAgent(ctx, einoAgenticChatModelAgentConfig{ + Name: id, + Description: desc, + Instruction: subInstrFinal, + GenModelInput: literalAgenticInstructionGenModelInput, + Model: subModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: subToolsForCfg, + UnknownToolsHandler: einomcp.UnknownToolReminderHandler(), + ToolCallMiddlewares: []compose.ToolMiddleware{ + modelOutputExecutionGuardMiddleware(), + localToolRBACMiddleware(), + hitlToolCallMiddleware(), + softRecoveryToolMiddleware(), + }, + }, + EmitInternalEvents: true, + }, + MaxIterations: subMax, + Handlers: subHandlers, + ModelRetryConfig: agenticModelRetryCfg, + ModelFailoverConfig: agenticModelFailoverCfg, + }) + if err != nil { + return nil, fmt.Errorf("子代理 %q: %w", id, err) + } + subAgents = append(subAgents, sa) + if adapted := newEinoAgenticMessageAgentAdapter(sa); adapted != nil { + supervisorSubAgents = append(supervisorSubAgents, adapted) + } + } + } + + modelFacingTrace := newModelFacingTraceHolder() + + // 与 deep.Config.Name / supervisor 主代理 Name 一致。 + orchestratorName := "cyberstrike-deep" + orchDescription := "Coordinates specialist agents and MCP tools for authorized security testing." + orchInstruction, orchMeta := resolveMainOrchestratorInstruction(orchMode, ma, markdownLoad) + if orchMeta != nil { + if strings.TrimSpace(orchMeta.EinoName) != "" { + orchestratorName = strings.TrimSpace(orchMeta.EinoName) + } + if d := strings.TrimSpace(orchMeta.Description); d != "" { + orchDescription = d + } + } else if orchMode == "deep" && orch != nil { + if strings.TrimSpace(orch.EinoName) != "" { + orchestratorName = strings.TrimSpace(orch.EinoName) + } + if d := strings.TrimSpace(orch.Description); d != "" { + orchDescription = d + } + } + + mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, orchestratorName) + if err != nil { + return nil, err + } + var mainToolsForCfg []tool.BaseTool + var mainToolSearchActive bool + var mainAgenticOrchestratorPre []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + mainToolsForCfg, mainAgenticOrchestratorPre, mainToolSearchActive, err = prependEinoAgenticMiddlewares(ctx, &ma.EinoMiddleware, einoMWMain, mainTools, agenticLoc, agenticSkillsRoot, conversationID, projectID, logger) + if err != nil { + return nil, err + } + + orchInstruction = project.AppendSystemPromptBlock(orchInstruction, systemPromptExtra) + orchInstruction = project.AppendVisionImageAnalysisIfReady(orchInstruction, appCfg.Vision.Ready()) + orchInstruction = injectToolNamesOnlyInstruction(ctx, orchInstruction, mainTools, mainToolSearchActive) + if logger != nil { + mainNames := collectToolNames(ctx, mainTools) + mountedNames := collectToolNames(ctx, mainToolsForCfg) + logger.Info("eino tool-name injection", + zap.String("scope", "orchestrator"), + zap.String("orchestration", orchMode), + zap.Int("tool_names", len(mainNames)), + zap.Int("mounted_tool_names", len(mountedNames)), + zap.Bool("tool_search_middleware", mainToolSearchActive), + ) + } + + supInstr := strings.TrimSpace(orchInstruction) + if orchMode == "supervisor" { + var sb strings.Builder + if supInstr != "" { + sb.WriteString(supInstr) + sb.WriteString("\n\n") + } + sb.WriteString("你是监督协调者:可将任务通过 transfer 工具委派给下列专家子代理(使用其在系统中的 Agent 名称)。专家列表:") + for _, sa := range subAgents { + if sa == nil { + continue + } + sb.WriteString("\n- ") + sb.WriteString(sa.Name(ctx)) + } + sb.WriteString("\n\nSupervisor 是专家路由模式:仅当任务确实需要不同专家分工时才 transfer;简单查询、单步工具调用或无需专业分流的任务由你直接完成。避免在同一子代理之间反复 transfer;除非有新的、具体的补充目标。专家返回后,你必须自行汇总、裁剪、校验证据,再用 exit 交付最终答案。") + sb.WriteString("\n\n当你已完成用户目标或需要将最终结论交付用户时,使用 exit 工具结束。") + supInstr = sb.String() + } + + var deepBackend filesystem.Backend + var deepShell filesystem.StreamingShell + if agenticLoc != nil && agenticFSTools { + deepBackend = agenticLoc + deepShell = &einoStreamingShellWrap{ + inner: security.NewEinoStreamingShell(), + invokeNotify: toolInvokeNotify, + einoAgentName: orchestratorName, + outputChunk: nil, + beginMonitor: einoExecBegin, + appendPartialMonitor: einoExecAppendPartial, + registerCancelMonitor: einoExecRegisterCancel, + unregisterCancelMonitor: einoExecUnregisterCancel, + finishMonitor: einoExecFinish, + toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg), + toolWaitTimeoutSeconds: agentToolWaitTimeoutSeconds(appCfg), + shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg), + } + } + + var mainModel model.AgenticModel + var mainSumMw adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + if orchMode != "plan_execute" { + mainModel, err = agenticModelFactory(ctx, appCfg.OpenAI, einoModelModeNormal) + if err != nil { + return nil, fmt.Errorf("多代理主 AgenticModel: %w", err) + } + mainSumMw, err = newEinoAgenticSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) + if err != nil { + return nil, fmt.Errorf("多代理主 agentic summarization 中间件: %w", err) + } + } + + // noNestedTaskMiddleware 必须在最外层(最先拦截),防止 skill 或其他中间件内部触发 task 调用绕过检测。 + deepHandlers := []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newNoNestedAgenticTaskMiddleware()} + var taskBlackboardSupplement string + if appCfg.Project.Enabled && db != nil { + if pid := strings.TrimSpace(projectID); pid != "" { + if block, err := project.BuildFactIndexBlock(db, pid, appCfg.Project); err == nil { + taskBlackboardSupplement = strings.TrimSpace(block) + } + } + } + if mw := newAgenticTaskContextEnrichMiddleware(runtimeUserMessage, history, ma.SubAgentUserContextMaxRunesEffective(), taskBlackboardSupplement); mw != nil { + deepHandlers = append(deepHandlers, mw) + } + if len(mainAgenticOrchestratorPre) > 0 { + deepHandlers = append(deepHandlers, mainAgenticOrchestratorPre...) + } + if agenticSkillMW != nil { + deepHandlers = append(deepHandlers, agenticSkillMW) + } + deepHandlers = appendEinoAgenticChatModelTailMiddlewares(deepHandlers, einoChatModelTailConfig{ + logger: logger, + phase: "deep_orchestrator", + agenticSummarization: mainSumMw, + modelName: appCfg.OpenAI.Model, + maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, + toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware), + conversationID: conversationID, + trace: modelFacingTrace, + middlewareConfig: &ma.EinoMiddleware, + }) + + supHandlers := []adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage]{} + if len(mainAgenticOrchestratorPre) > 0 { + supHandlers = append(supHandlers, mainAgenticOrchestratorPre...) + } + if agenticSkillMW != nil { + supHandlers = append(supHandlers, agenticSkillMW) + } + supHandlers = appendEinoAgenticChatModelTailMiddlewares(supHandlers, einoChatModelTailConfig{ + logger: logger, + phase: "supervisor_orchestrator", + agenticSummarization: mainSumMw, + modelName: appCfg.OpenAI.Model, + maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, + toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware), + conversationID: conversationID, + trace: modelFacingTrace, + middlewareConfig: &ma.EinoMiddleware, + }) + + mainToolsCfg := adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: mainToolsForCfg, + UnknownToolsHandler: einomcp.UnknownToolReminderHandler(), + ToolCallMiddlewares: []compose.ToolMiddleware{ + modelOutputExecutionGuardMiddleware(), + localToolRBACMiddleware(), + hitlToolCallMiddleware(), + softRecoveryToolMiddleware(), + }, + }, + EmitInternalEvents: true, + } + + deepAgenticOutKey, agenticTaskGen := deepAgenticExtrasFromConfig(ma) + + var da adk.Agent + switch orchMode { + case "plan_execute": + peMainModel, perr := modelFactory(ctx, appCfg.OpenAI, einoModelModePlanner) + if perr != nil { + return nil, fmt.Errorf("plan_execute 规划模型: %w", perr) + } + if logger != nil { + logger.Info("plan_execute: planner/replanner 使用无 reasoning 的独立 ChatModel(ToolChoiceForced 兼容)", + zap.String("model", appCfg.OpenAI.Model), + ) + } + execModel, perr := modelFactory(ctx, appCfg.OpenAI, einoModelModeNormal) + if perr != nil { + return nil, fmt.Errorf("plan_execute 执行器模型: %w", perr) + } + agenticExecModel, perr := agenticModelFactory(ctx, appCfg.OpenAI, einoModelModeNormal) + if perr != nil { + return nil, fmt.Errorf("plan_execute 执行器 AgenticModel: %w", perr) + } + planRewriteSumMw, perr := newEinoSummarizationMiddleware(ctx, execModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger) + if perr != nil { + return nil, fmt.Errorf("plan_execute planner/replanner summarization: %w", perr) + } + var peFsMw adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] + if agenticSkillMW != nil && agenticFSTools && agenticLoc != nil { + peFsMw, err = subAgentAgenticFilesystemMiddleware(ctx, agenticLoc, toolInvokeNotify, "executor", conversationID, projectID, ma.EinoMiddleware.ReductionRootDir, toolMaxBytesFromMW(&ma.EinoMiddleware), mcpExecBinder, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil) + if err != nil { + return nil, fmt.Errorf("plan_execute agentic filesystem 中间件: %w", err) + } + } + peRoot, perr := NewPlanExecuteRoot(ctx, &PlanExecuteRootArgs{ + MainToolCallingModel: peMainModel, + AgenticExecModel: agenticExecModel, + OrchInstruction: orchInstruction, + ToolsCfg: mainToolsCfg, + ExecMaxIter: deepMaxIter, + LoopMaxIter: ma.PlanExecuteLoopMaxIterations, + AppCfg: appCfg, + MwCfg: &ma.EinoMiddleware, + ConversationID: conversationID, + DB: db, + ProjectID: projectID, + Logger: logger, + ModelName: appCfg.OpenAI.Model, + // 与 Deep/Supervisor 主代理同源:typed patch / reduction / toolsearch / plantask(见 buildPlanExecuteAgenticExecutorHandlers)。 + AgenticExecPreMiddlewares: mainAgenticOrchestratorPre, + AgenticSkillMiddleware: agenticSkillMW, + AgenticFilesystemMiddleware: peFsMw, + ModelFacingTrace: modelFacingTrace, + PlannerReplannerRewriteHandlers: appendEinoChatModelTailMiddlewares(nil, einoChatModelTailConfig{ + logger: logger, + phase: "plan_execute_planner_replanner", + summarization: planRewriteSumMw, + modelName: appCfg.OpenAI.Model, + maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, + toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware), + conversationID: conversationID, + skipTrace: true, + middlewareConfig: &ma.EinoMiddleware, + }), + AgenticModelRetryConfig: agenticModelRetryCfg, + AgenticModelFailoverConfig: agenticModelFailoverCfg, + }) + if perr != nil { + return nil, perr + } + da = peRoot + case "supervisor": + supCfg := einoAgenticChatModelAgentConfig{ + Name: orchestratorName, + Description: orchDescription, + Instruction: supInstr, + GenModelInput: literalAgenticInstructionGenModelInput, + Model: mainModel, + ToolsConfig: mainToolsCfg, + MaxIterations: deepMaxIter, + Handlers: supHandlers, + Exit: &adk.ExitTool{}, + ModelRetryConfig: agenticModelRetryCfg, + ModelFailoverConfig: agenticModelFailoverCfg, + } + if deepAgenticOutKey != "" { + supCfg.OutputKey = deepAgenticOutKey + } + superChat, serr := newEinoAgenticChatModelAgentAdapter(ctx, supCfg) + if serr != nil { + return nil, fmt.Errorf("supervisor agentic 主代理: %w", serr) + } + supRoot, serr := supervisor.New(ctx, &supervisor.Config{ + Supervisor: superChat, + SubAgents: supervisorSubAgents, + }) + if serr != nil { + return nil, fmt.Errorf("supervisor.New: %w", serr) + } + da = supRoot + default: + dcfg := &deep.TypedConfig[*schema.AgenticMessage]{ + Name: orchestratorName, + Description: orchDescription, + ChatModel: mainModel, + Instruction: orchInstruction, + SubAgents: subAgents, + WithoutGeneralSubAgent: ma.WithoutGeneralSubAgent, + WithoutWriteTodos: ma.WithoutWriteTodos, + MaxIteration: deepMaxIter, + Backend: deepBackend, + StreamingShell: deepShell, + Handlers: deepHandlers, + ToolsConfig: mainToolsCfg, + ModelRetryConfig: agenticModelRetryCfg, + ModelFailoverConfig: agenticModelFailoverCfg, + } + if deepAgenticOutKey != "" { + dcfg.OutputKey = deepAgenticOutKey + } + if agenticTaskGen != nil { + dcfg.TaskToolDescriptionGenerator = agenticTaskGen + } + dDeep, derr := deep.NewTyped[*schema.AgenticMessage](ctx, dcfg) + if derr != nil { + return nil, fmt.Errorf("deep.NewTyped[AgenticMessage]: %w", derr) + } + da = newEinoAgenticMessageAgentAdapter(dDeep) + } + + baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware) + baseMsgs = appendUserMessageIfNeeded(baseMsgs, runtimeUserMessage) + + streamsMainAssistant := func(agent string) bool { + if orchMode == "plan_execute" { + return planExecuteStreamsMainAssistant(agent) + } + return agent == "" || agent == orchestratorName + } + einoRoleTag := func(agent string) string { + if orchMode == "plan_execute" { + return planExecuteEinoRoleTag(agent) + } + if streamsMainAssistant(agent) { + return "orchestrator" + } + return "sub" + } + + return runEinoADKAgentLoop(ctx, &einoADKRunLoopArgs{ + OrchMode: orchMode, + OrchestratorName: orchestratorName, + ConversationID: conversationID, + Progress: progress, + Logger: logger, + SnapshotMCPIDs: snapshotMCPIDs, + StreamsMainAssistant: streamsMainAssistant, + EinoRoleTag: einoRoleTag, + CheckpointDir: ma.EinoMiddleware.CheckpointDir, + RunRetryMaxAttempts: RunRetryMaxAttemptsFromConfig(&ma.EinoMiddleware), + RunRetryMaxBackoffSec: int(einoRunRetryMaxBackoffFromConfig(&ma.EinoMiddleware).Seconds()), + McpIDsMu: &mcpIDsMu, + McpIDs: &mcpIDs, + FilesystemMonitorAgent: ag, + FilesystemMonitorRecord: recorder, + MCPExecutionBinder: mcpExecBinder, + ToolInvokeNotify: toolInvokeNotify, + DA: da, + ModelFacingTrace: modelFacingTrace, + EinoCallbacks: &ma.EinoCallbacks, + MaxTotalTokens: appCfg.OpenAI.MaxTotalTokens, + ToolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware), + ModelName: appCfg.OpenAI.Model, + MiddlewareConfig: &ma.EinoMiddleware, + EmptyResponseMessage: "(Eino multi-agent orchestration completed but no assistant text was captured. Check process details or logs.) " + + "(Eino 多代理编排已完成,但未捕获到助手文本输出。请查看过程详情或日志。)", + }, baseMsgs) +} + +func chatToolCallsToSchema(tcs []agent.ToolCall) []schema.ToolCall { + if len(tcs) == 0 { + return nil + } + out := make([]schema.ToolCall, 0, len(tcs)) + for _, tc := range tcs { + if strings.TrimSpace(tc.ID) == "" { + continue + } + argsStr := "" + if tc.Function.Arguments != nil { + b, err := json.Marshal(tc.Function.Arguments) + if err == nil { + argsStr = string(b) + } + } + // Some OpenAI-compatible gateways require `function.arguments` to exist + // on every assistant tool_call message. When args are empty, omitempty may + // drop the field during serialization and cause "missing field arguments" + // on the next turn history replay. + if strings.TrimSpace(argsStr) == "" { + argsStr = "{}" + } + typ := tc.Type + if typ == "" { + typ = "function" + } + out = append(out, schema.ToolCall{ + ID: tc.ID, + Type: typ, + Function: schema.FunctionCall{ + Name: tc.Function.Name, + Arguments: argsStr, + }, + }) + } + return out +} + +// historyToMessages 将已保存的 model-facing 轨迹转为 Eino ADK 消息。 +// 新轨迹应已是模型实际看到的内容;对旧版本遗留的超大 tool 正文再做一次上限规范化, +// 防止原始工具输出通过 last_react/checkpoint 绕过 reduction。 +func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message { + toolContentMax := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective() + userContentMaxRunes := config.MultiAgentEinoMiddlewareConfig{}.LatestUserMessageMaxRunesEffective() + if mwCfg != nil { + toolContentMax = mwCfg.ReductionMaxLengthForTruncEffective() + userContentMaxRunes = mwCfg.LatestUserMessageMaxRunesEffective() + } + if appCfg != nil { + userContentMaxRunes = minPositiveInt(userContentMaxRunes, modelFacingRuneBudget(appCfg.OpenAI.MaxTotalTokens, 0.20)) + } + if len(history) == 0 { + return nil + } + raw := make([]adk.Message, 0, len(history)) + for _, h := range history { + role := strings.ToLower(strings.TrimSpace(h.Role)) + switch role { + case "user": + if strings.TrimSpace(h.Content) != "" { + content := h.Content + if !h.ModelFacingTrace { + content = normalizeRestoredUserContent(content, userContentMaxRunes) + } + raw = append(raw, schema.UserMessage(content)) + } + case "assistant": + toolSchema := chatToolCallsToSchema(h.ToolCalls) + hasRC := strings.TrimSpace(h.ReasoningContent) != "" + if len(toolSchema) > 0 || strings.TrimSpace(h.Content) != "" || hasRC { + am := schema.AssistantMessage(h.Content, toolSchema) + if hasRC { + am.ReasoningContent = strings.TrimSpace(h.ReasoningContent) + } + raw = append(raw, am) + } + case "tool": + if strings.TrimSpace(h.ToolCallID) == "" && strings.TrimSpace(h.Content) == "" { + continue + } + var opts []schema.ToolMessageOption + if tn := strings.TrimSpace(h.ToolName); tn != "" { + opts = append(opts, schema.WithToolName(tn)) + } + content := h.Content + if !h.ModelFacingTrace || (toolContentMax > 0 && len(content) > toolContentMax) { + content = normalizeRestoredToolContent(content, toolContentMax) + } + raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...)) + default: + continue + } + } + return raw +} + +func normalizeRestoredUserContent(content string, maxRunes int) string { + if maxRunes <= 0 || utf8.RuneCountInString(content) <= maxRunes { + return content + } + runes := []rune(content) + const marker = "\n\n...[historical user input normalized to the model-facing budget]...\n\n" + markerRunes := []rune(marker) + budget := maxRunes - len(markerRunes) + if budget <= 0 { + end := maxRunes + if end > len(markerRunes) { + end = len(markerRunes) + } + return string(markerRunes[:end]) + } + head := budget / 2 + tail := budget - head + return string(runes[:head]) + marker + string(runes[len(runes)-tail:]) +} + +func normalizeRestoredToolContent(content string, maxBytes int) string { + if maxBytes <= 0 || len(content) <= maxBytes { + return content + } + const marker = "\n\n...[legacy tool output discarded during model-facing history migration]...\n\n" + budget := maxBytes - len(marker) + if budget <= 0 { + 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:] +} + +// mergeStreamingToolCallFragments 将流式多帧的 ToolCall 按 index 合并 arguments(与 schema.concatToolCalls 行为一致)。 +func mergeStreamingToolCallFragments(fragments []schema.ToolCall) []schema.ToolCall { + if len(fragments) == 0 { + return nil + } + m, err := schema.ConcatMessages([]*schema.Message{{ToolCalls: fragments}}) + if err != nil || m == nil { + return fragments + } + return m.ToolCalls +} + +// mergeMessageToolCalls 非流式路径上若仍带分片式 tool_calls,合并后再上报 UI。 +func mergeMessageToolCalls(msg *schema.Message) *schema.Message { + if msg == nil || len(msg.ToolCalls) == 0 { + return msg + } + m, err := schema.ConcatMessages([]*schema.Message{msg}) + if err != nil || m == nil { + return msg + } + out := *msg + out.ToolCalls = m.ToolCalls + return &out +} + +// toolCallStableID 用于流式阶段去重;OpenAI 流式常先给 index 后补 id。 +func toolCallStableID(tc schema.ToolCall) string { + if tc.ID != "" { + return tc.ID + } + if tc.Index != nil { + return fmt.Sprintf("idx:%d", *tc.Index) + } + return "" +} + +// toolCallDisplayName returns the visible tool name once the model stream has +// produced a concrete function name. Anonymous stream fragments are filtered +// before progress emission instead of being guessed as task calls. +func toolCallDisplayName(tc schema.ToolCall) string { + if n := strings.TrimSpace(tc.Function.Name); n != "" { + return n + } + if n := strings.TrimSpace(tc.Type); n != "" && !strings.EqualFold(n, "function") { + return n + } + return "" +} + +// toolCallsSignatureFlush 用于去重键;无 id/index 时用占位 pos,避免流末帧缺 id 时整条工具事件丢失。 +func toolCallsSignatureFlush(msg *schema.Message) string { + if msg == nil || len(msg.ToolCalls) == 0 { + return "" + } + visible := filterVisibleToolCallsForProgress(msg.ToolCalls) + if len(visible) == 0 { + return "" + } + parts := make([]string, 0, len(visible)) + for i, tc := range visible { + id := toolCallStableID(tc) + if id == "" { + id = fmt.Sprintf("pos:%d", i) + } + name := toolCallDisplayName(tc) + if name == "" { + continue + } + parts = append(parts, id+"|"+name) + } + if len(parts) == 0 { + return "" + } + sort.Strings(parts) + return strings.Join(parts, ";") +} + +// toolCallsRichSignature 用于去重:同一次流式已上报后,紧随其后的非流式消息常带相同 tool_calls。 +func toolCallsRichSignature(msg *schema.Message) string { + base := toolCallsSignatureFlush(msg) + if base == "" { + return "" + } + visible := filterVisibleToolCallsForProgress(msg.ToolCalls) + parts := make([]string, 0, len(visible)) + for _, tc := range visible { + id := toolCallStableID(tc) + arg := tc.Function.Arguments + if len(arg) > 240 { + arg = arg[:240] + } + parts = append(parts, id+":"+arg) + } + sort.Strings(parts) + return base + "|" + strings.Join(parts, ";") +} + +func einoMainIterationKey(agentName, orchestratorName string) string { + key := strings.TrimSpace(agentName) + if key == "" { + key = strings.TrimSpace(orchestratorName) + } + if key == "" { + return "_main" + } + return key +} + +func tryEmitToolCallsOnce( + msg *schema.Message, + agentName, orchestratorName, conversationID, orchMode string, + progress func(string, string, interface{}), + seen map[string]struct{}, + subAgentToolStep, mainAgentToolStep map[string]int, + markPending func(toolCallPendingInfo), +) { + if msg == nil || len(msg.ToolCalls) == 0 || progress == nil || seen == nil { + return + } + if toolCallsSignatureFlush(msg) == "" { + return + } + sig := agentName + "\x1e" + toolCallsRichSignature(msg) + if _, ok := seen[sig]; ok { + return + } + seen[sig] = struct{}{} + emitToolCallsFromMessage(msg, agentName, orchestratorName, conversationID, orchMode, progress, subAgentToolStep, mainAgentToolStep, markPending) +} + +func emitToolCallsFromMessage( + msg *schema.Message, + agentName, orchestratorName, conversationID, orchMode string, + progress func(string, string, interface{}), + subAgentToolStep, mainAgentToolStep map[string]int, + markPending func(toolCallPendingInfo), +) { + if msg == nil || len(msg.ToolCalls) == 0 || progress == nil { + return + } + visibleToolCalls := filterVisibleToolCallsForProgress(msg.ToolCalls) + if len(visibleToolCalls) == 0 { + return + } + if subAgentToolStep == nil { + subAgentToolStep = make(map[string]int) + } + isSubToolRound := agentName != "" && agentName != orchestratorName + if isSubToolRound { + subAgentToolStep[agentName]++ + n := subAgentToolStep[agentName] + progress("iteration", "", map[string]interface{}{ + "iteration": n, + "einoScope": "sub", + "einoRole": "sub", + "einoAgent": agentName, + "conversationId": conversationID, + "source": "eino", + }) + } else if mainAgentToolStep != nil { + key := einoMainIterationKey(agentName, orchestratorName) + mainAgentToolStep[key]++ + n := mainAgentToolStep[key] + // 第 1 轮已在主代理进入时发出;此后每次工具批次对应新一轮 ReAct(与子代理按工具计步一致)。 + if n > 1 { + progress("iteration", "", map[string]interface{}{ + "iteration": n, + "einoScope": "main", + "einoRole": "orchestrator", + "einoAgent": agentName, + "orchestration": orchMode, + "conversationId": conversationID, + "source": "eino", + }) + } + } + role := "orchestrator" + if isSubToolRound { + role = "sub" + } + progress("tool_calls_detected", fmt.Sprintf("检测到 %d 个工具调用", len(visibleToolCalls)), map[string]interface{}{ + "count": len(visibleToolCalls), + "conversationId": conversationID, + "source": "eino", + "einoAgent": agentName, + "einoRole": role, + }) + for idx, tc := range visibleToolCalls { + argStr := strings.TrimSpace(tc.Function.Arguments) + if argStr == "" && len(tc.Extra) > 0 { + if b, mErr := json.Marshal(tc.Extra); mErr == nil { + argStr = string(b) + } + } + var argsObj map[string]interface{} + if argStr != "" { + if uErr := json.Unmarshal([]byte(argStr), &argsObj); uErr != nil || argsObj == nil { + argsObj = map[string]interface{}{"_raw": argStr} + } + } + display := toolCallDisplayName(tc) + toolCallID := tc.ID + if toolCallID == "" && tc.Index != nil { + // Stream indexes restart from zero for every model turn. Include a + // process-wide sequence so pending/result de-duplication cannot collide + // with an earlier batch in the same agent run. + toolCallID = fmt.Sprintf("eino-stream-%d-%d", fallbackToolCallSequence.Add(1), *tc.Index) + } + // Record visible pending tool calls for later tool_result correlation / recovery flushing. + if markPending != nil && toolCallID != "" { + markPending(toolCallPendingInfo{ + ToolCallID: toolCallID, + ToolName: display, + Arguments: argsObj, + EinoAgent: agentName, + EinoRole: role, + }) + } + progress("tool_call", fmt.Sprintf("正在调用工具: %s", display), map[string]interface{}{ + "toolName": display, + "arguments": argStr, + "argumentsObj": argsObj, + "toolCallId": toolCallID, + "index": idx + 1, + "total": len(visibleToolCalls), + "conversationId": conversationID, + "source": "eino", + "einoAgent": agentName, + "einoRole": role, + }) + } +} + +func filterVisibleToolCallsForProgress(calls []schema.ToolCall) []schema.ToolCall { + if len(calls) == 0 { + return nil + } + out := make([]schema.ToolCall, 0, len(calls)) + for _, tc := range calls { + if _, ok := modelOutputRecoveryFromToolCall(tc); ok { + continue + } + if toolCallDisplayName(tc) == "" { + continue + } + out = append(out, tc) + } + return out +} + +// dedupeRepeatedParagraphs 去掉完全相同的连续/重复段落,缓解多代理各自复述同一列表。 +func dedupeRepeatedParagraphs(s string, minLen int) string { + if s == "" || minLen <= 0 { + return s + } + paras := strings.Split(s, "\n\n") + var out []string + seen := make(map[string]bool) + for _, p := range paras { + t := strings.TrimSpace(p) + if len(t) < minLen { + out = append(out, p) + continue + } + if seen[t] { + continue + } + seen[t] = true + out = append(out, p) + } + return strings.TrimSpace(strings.Join(out, "\n\n")) +} + +// dedupeParagraphsByLineFingerprint 去掉「正文行集合相同」的重复段落(开场白略不同也会合并),缓解多代理各写一遍目录清单。 +func dedupeParagraphsByLineFingerprint(s string, minParaLen int) string { + if s == "" || minParaLen <= 0 { + return s + } + paras := strings.Split(s, "\n\n") + var out []string + seen := make(map[string]bool) + for _, p := range paras { + t := strings.TrimSpace(p) + if len(t) < minParaLen { + out = append(out, p) + continue + } + fp := paragraphLineFingerprint(t) + // 指纹仅在「≥4 条非空行」时有效;单行/短段落长回复(如自我介绍)fp 为空,必须保留,否则会误删全文并触发「未捕获到助手文本」占位。 + if fp == "" { + out = append(out, p) + continue + } + if seen[fp] { + continue + } + seen[fp] = true + out = append(out, p) + } + return strings.TrimSpace(strings.Join(out, "\n\n")) +} + +func paragraphLineFingerprint(t string) string { + lines := strings.Split(t, "\n") + norm := make([]string, 0, len(lines)) + for _, L := range lines { + s := strings.TrimSpace(L) + if s == "" { + continue + } + norm = append(norm, s) + } + if len(norm) < 4 { + return "" + } + sort.Strings(norm) + return strings.Join(norm, "\x1e") +} diff --git a/internal/multiagent/runner_tool_call_id_test.go b/internal/multiagent/runner_tool_call_id_test.go new file mode 100644 index 00000000..e3e89356 --- /dev/null +++ b/internal/multiagent/runner_tool_call_id_test.go @@ -0,0 +1,37 @@ +package multiagent + +import ( + "testing" + + "github.com/cloudwego/eino/schema" +) + +func TestEmitToolCallsUsesUniqueFallbackIDsAcrossBatches(t *testing.T) { + index := 0 + msg := &schema.Message{ToolCalls: []schema.ToolCall{{ + Index: &index, + Function: schema.FunctionCall{ + Name: "http-framework-test", + Arguments: `{}`, + }, + }}} + var ids []string + progress := func(eventType, _ string, raw interface{}) { + if eventType != "tool_call" { + return + } + data, _ := raw.(map[string]interface{}) + if id, _ := data["toolCallId"].(string); id != "" { + ids = append(ids, id) + } + } + for i := 0; i < 2; i++ { + emitToolCallsFromMessage(msg, "agent", "agent", "conversation", "deep", progress, nil, make(map[string]int), nil) + } + if len(ids) != 2 { + t.Fatalf("fallback IDs = %v, want two IDs", ids) + } + if ids[0] == ids[1] { + t.Fatalf("fallback ID was reused across batches: %q", ids[0]) + } +} diff --git a/internal/multiagent/shell_tool_guidance.go b/internal/multiagent/shell_tool_guidance.go new file mode 100644 index 00000000..0eec25bf --- /dev/null +++ b/internal/multiagent/shell_tool_guidance.go @@ -0,0 +1,33 @@ +package multiagent + +import ( + "strings" + + "cyberstrike-ai/internal/projectprompt" +) + +func shellToolsPresent(toolNames []string) bool { + for _, n := range toolNames { + switch strings.ToLower(strings.TrimSpace(n)) { + case "exec", "execute": + return true + } + } + return false +} + +// injectShellToolGuidance 在系统提示末尾追加 exec/execute 分工(仅当工具列表含 exec 或 execute)。 +func injectShellToolGuidance(instruction string, toolNames []string) string { + if !shellToolsPresent(toolNames) { + return instruction + } + block := strings.TrimSpace(projectprompt.ShellExecExecuteGuidanceSection()) + if block == "" { + return instruction + } + s := strings.TrimSpace(instruction) + if s == "" { + return block + } + return s + "\n\n" + block +} diff --git a/internal/multiagent/shell_tool_guidance_test.go b/internal/multiagent/shell_tool_guidance_test.go new file mode 100644 index 00000000..91bd2730 --- /dev/null +++ b/internal/multiagent/shell_tool_guidance_test.go @@ -0,0 +1,17 @@ +package multiagent + +import ( + "strings" + "testing" +) + +func TestInjectShellToolGuidance(t *testing.T) { + got := injectShellToolGuidance("base", []string{"nmap"}) + if got != "base" { + t.Fatalf("expected unchanged, got %q", got) + } + got = injectShellToolGuidance("base", []string{"exec", "nmap"}) + if !strings.Contains(got, "exec/execute") || !strings.Contains(got, "base") { + t.Fatalf("expected shell guidance appended, got %q", got) + } +} diff --git a/internal/multiagent/sub_agent_context.go b/internal/multiagent/sub_agent_context.go new file mode 100644 index 00000000..d7fba52b --- /dev/null +++ b/internal/multiagent/sub_agent_context.go @@ -0,0 +1,208 @@ +package multiagent + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "cyberstrike-ai/internal/agent" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +const userContextSupplementHeader = "\n\n## 用户历史输入(原文,子代理必读)\n" + +// taskContextEnrichMiddleware intercepts "task" tool calls on the orchestrator +// and appends the user's original conversation messages to the task description. +// This ensures sub-agents always receive the full user intent (target URLs, +// scope, etc.) even when the orchestrator forgets to include them. +// +// Design: user context is injected into the task description (per-task), NOT +// into the sub-agent's Instruction (system prompt). This keeps sub-agent +// Instructions clean as pure role definitions while attaching context to the +// specific delegation — aligned with Claude Code's agent design philosophy. +type taskContextEnrichMiddleware struct { + adk.BaseChatModelAgentMiddleware + supplement string // pre-built user context block +} + +// newTaskContextEnrichMiddleware returns a middleware that enriches task +// descriptions with user conversation context. Returns nil if disabled +// (maxRunes < 0) or no user messages exist. +// projectBlackboard 仅传项目黑板索引块(BuildFactIndexBlock);勿传完整 systemPromptExtra。 +func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int, projectBlackboard string) adk.ChatModelAgentMiddleware { + supplement := buildUserContextSupplement(userMessage, history, maxRunes) + if bb := strings.TrimSpace(projectBlackboard); bb != "" { + if supplement != "" { + supplement += "\n\n" + bb + } else { + supplement = "\n\n" + bb + } + } + if supplement == "" { + return nil + } + return &taskContextEnrichMiddleware{supplement: supplement} +} + +type agenticTaskContextEnrichMiddleware struct { + *adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + supplement string +} + +func newAgenticTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int, projectBlackboard string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + supplement := buildUserContextSupplement(userMessage, history, maxRunes) + if bb := strings.TrimSpace(projectBlackboard); bb != "" { + if supplement != "" { + supplement += "\n\n" + bb + } else { + supplement = "\n\n" + bb + } + } + if supplement == "" { + return nil + } + return &agenticTaskContextEnrichMiddleware{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + supplement: supplement, + } +} + +func (m *taskContextEnrichMiddleware) WrapInvokableToolCall( + ctx context.Context, + endpoint adk.InvokableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.InvokableToolCallEndpoint, error) { + return wrapTaskContextEnrichCall(m, ctx, endpoint, tCtx) +} + +func (m *agenticTaskContextEnrichMiddleware) WrapInvokableToolCall( + ctx context.Context, + endpoint adk.InvokableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.InvokableToolCallEndpoint, error) { + return wrapTaskContextEnrichCall(m, ctx, endpoint, tCtx) +} + +type taskContextEnricher interface { + enrichTaskDescription(argsJSON string) string +} + +func wrapTaskContextEnrichCall( + m taskContextEnricher, + ctx context.Context, + endpoint adk.InvokableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.InvokableToolCallEndpoint, error) { + if tCtx == nil || !strings.EqualFold(strings.TrimSpace(tCtx.Name), "task") { + return endpoint, nil + } + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + enriched := m.enrichTaskDescription(argumentsInJSON) + return endpoint(ctx, enriched, opts...) + }, nil +} + +// enrichTaskDescription parses the task JSON arguments, appends user context +// to the "description" field, and re-serializes. Falls back to the original +// JSON if parsing fails or no description field exists. +func (m *taskContextEnrichMiddleware) enrichTaskDescription(argsJSON string) string { + return enrichTaskDescriptionWithSupplement(argsJSON, m.supplement) +} + +func (m *agenticTaskContextEnrichMiddleware) enrichTaskDescription(argsJSON string) string { + return enrichTaskDescriptionWithSupplement(argsJSON, m.supplement) +} + +func enrichTaskDescriptionWithSupplement(argsJSON, supplement string) string { + var raw map[string]interface{} + if err := json.Unmarshal([]byte(argsJSON), &raw); err != nil { + return argsJSON + } + desc, ok := raw["description"].(string) + if !ok { + return argsJSON + } + raw["description"] = desc + supplement + enriched, err := json.Marshal(raw) + if err != nil { + return argsJSON + } + return string(enriched) +} + +// buildUserContextSupplement collects user messages from conversation history +// and the current message, returning a formatted block to append to task +// descriptions. Returns "" if disabled or no user messages exist. +func buildUserContextSupplement(userMessage string, history []agent.ChatMessage, maxRunes int) string { + if maxRunes < 0 { + return "" + } + + var userMsgs []string + for _, h := range history { + if h.Role == "user" { + if m := strings.TrimSpace(h.Content); m != "" { + userMsgs = append(userMsgs, m) + } + } + } + if um := strings.TrimSpace(userMessage); um != "" { + if len(userMsgs) == 0 || userMsgs[len(userMsgs)-1] != um { + userMsgs = append(userMsgs, um) + } + } + if len(userMsgs) == 0 { + return "" + } + + lines := make([]string, 0, len(userMsgs)) + for i, msg := range userMsgs { + lines = append(lines, fmt.Sprintf("[第%d轮] %s", i+1, msg)) + } + joined := strings.Join(lines, "\n") + if maxRunes > 0 && len([]rune(joined)) > maxRunes { + joined = truncateKeepFirstLast(userMsgs, maxRunes) + } + + return userContextSupplementHeader + joined +} + +// truncateKeepFirstLast keeps the first and last user messages, giving each +// half the rune budget. The first message typically contains target info; +// the last contains the current instruction. +func truncateKeepFirstLast(msgs []string, maxRunes int) string { + if len(msgs) == 1 { + return truncateRunes(msgs[0], maxRunes) + } + + first := msgs[0] + last := msgs[len(msgs)-1] + sep := "\n---\n...(中间对话省略)...\n---\n" + sepLen := len([]rune(sep)) + + budget := maxRunes - sepLen + if budget <= 0 { + return truncateRunes(first+"\n---\n"+last, maxRunes) + } + + halfBudget := budget / 2 + firstTrunc := truncateRunes(first, halfBudget) + lastTrunc := truncateRunes(last, budget-len([]rune(firstTrunc))) + + return firstTrunc + sep + lastTrunc +} + +func truncateRunes(s string, max int) string { + rs := []rune(s) + if len(rs) <= max { + return s + } + if max <= 0 { + return "" + } + return string(rs[:max]) +} diff --git a/internal/multiagent/sub_agent_context_test.go b/internal/multiagent/sub_agent_context_test.go new file mode 100644 index 00000000..e5e030a5 --- /dev/null +++ b/internal/multiagent/sub_agent_context_test.go @@ -0,0 +1,183 @@ +package multiagent + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "cyberstrike-ai/internal/agent" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" +) + +// --- buildUserContextSupplement tests --- + +func TestBuildUserContextSupplement_SingleMessage(t *testing.T) { + result := buildUserContextSupplement("http://8.163.32.73:8081 测试命令执行", nil, 0) + if result == "" { + t.Fatal("expected non-empty supplement") + } + if !strings.Contains(result, "http://8.163.32.73:8081") { + t.Error("expected URL in supplement") + } +} + +func TestBuildUserContextSupplement_MultiTurn(t *testing.T) { + history := []agent.ChatMessage{ + {Role: "user", Content: "http://8.163.32.73:8081 这是一个pikachu靶场,尝试测试命令执行"}, + {Role: "assistant", Content: "好的,我来测试..."}, + {Role: "user", Content: "继续,并持久化webshell"}, + {Role: "assistant", Content: "正在处理..."}, + } + result := buildUserContextSupplement("你好", history, 0) + if !strings.Contains(result, "http://8.163.32.73:8081") { + t.Error("expected first turn URL to be preserved") + } + if !strings.Contains(result, "你好") { + t.Error("expected current message") + } +} + +func TestBuildUserContextSupplement_Empty(t *testing.T) { + if result := buildUserContextSupplement("", nil, 0); result != "" { + t.Errorf("expected empty, got %q", result) + } +} + +func TestBuildUserContextSupplement_Deduplicate(t *testing.T) { + history := []agent.ChatMessage{{Role: "user", Content: "你好"}} + result := buildUserContextSupplement("你好", history, 0) + if strings.Count(result, "你好") != 1 { + t.Errorf("expected '你好' once, got: %s", result) + } +} + +func TestBuildUserContextSupplement_SkipsNonUser(t *testing.T) { + history := []agent.ChatMessage{ + {Role: "user", Content: "目标是 10.0.0.1"}, + {Role: "assistant", Content: "不应该出现"}, + } + result := buildUserContextSupplement("确认", history, 0) + if strings.Contains(result, "不应该出现") { + t.Error("assistant message should not be included") + } +} + +func TestBuildUserContextSupplement_DisabledByNegative(t *testing.T) { + if result := buildUserContextSupplement("test", nil, -1); result != "" { + t.Errorf("expected empty when disabled, got %q", result) + } +} + +func TestBuildUserContextSupplement_CustomMaxRunes(t *testing.T) { + msg := strings.Repeat("A", 200) + result := buildUserContextSupplement(msg, nil, 50) + header := userContextSupplementHeader + body := strings.TrimPrefix(result, header) + if len([]rune(body)) > 50 { + t.Errorf("body should be capped at 50 runes, got %d", len([]rune(body))) + } +} + +func TestBuildUserContextSupplement_TruncateKeepsFirstAndLast(t *testing.T) { + first := "http://target.com " + strings.Repeat("A", 500) + var history []agent.ChatMessage + history = append(history, agent.ChatMessage{Role: "user", Content: first}) + for i := 0; i < 10; i++ { + history = append(history, agent.ChatMessage{Role: "user", Content: strings.Repeat("B", 500)}) + } + last := "最后一条指令" + result := buildUserContextSupplement(last, history, 800) + if !strings.Contains(result, "http://target.com") { + t.Error("first message (target URL) should survive truncation") + } + if !strings.Contains(result, last) { + t.Error("last message should survive truncation") + } +} + +// --- middleware integration tests --- + +func TestTaskContextEnrichMiddleware_EnrichesTaskDescription(t *testing.T) { + mw := newTaskContextEnrichMiddleware( + "继续测试", + []agent.ChatMessage{{Role: "user", Content: "http://8.163.32.73:8081 pikachu靶场"}}, + 0, + "", + ) + if mw == nil { + t.Fatal("expected non-nil middleware") + } + + called := false + var capturedArgs string + fakeEndpoint := func(ctx context.Context, args string, opts ...tool.Option) (string, error) { + called = true + capturedArgs = args + return "ok", nil + } + + wrapped, err := mw.(interface { + WrapInvokableToolCall(context.Context, adk.InvokableToolCallEndpoint, *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) + }).WrapInvokableToolCall(context.Background(), fakeEndpoint, &adk.ToolContext{Name: "task"}) + if err != nil { + t.Fatal(err) + } + + taskArgs := `{"subagent_type":"recon","description":"扫描目标端口"}` + wrapped(context.Background(), taskArgs) + + if !called { + t.Fatal("endpoint was not called") + } + + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(capturedArgs), &parsed); err != nil { + t.Fatalf("enriched args not valid JSON: %v", err) + } + desc := parsed["description"].(string) + if !strings.Contains(desc, "扫描目标端口") { + t.Error("original description should be preserved") + } + if !strings.Contains(desc, "http://8.163.32.73:8081") { + t.Error("user context should be appended to description") + } + if !strings.Contains(desc, "继续测试") { + t.Error("current user message should be in description") + } +} + +func TestTaskContextEnrichMiddleware_IgnoresNonTaskTools(t *testing.T) { + mw := newTaskContextEnrichMiddleware("test", nil, 0, "") + if mw == nil { + t.Fatal("expected non-nil middleware") + } + + original := `{"command":"nmap -sV target"}` + var capturedArgs string + fakeEndpoint := func(ctx context.Context, args string, opts ...tool.Option) (string, error) { + capturedArgs = args + return "ok", nil + } + + wrapped, err := mw.(interface { + WrapInvokableToolCall(context.Context, adk.InvokableToolCallEndpoint, *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) + }).WrapInvokableToolCall(context.Background(), fakeEndpoint, &adk.ToolContext{Name: "nmap_scan"}) + if err != nil { + t.Fatal(err) + } + + wrapped(context.Background(), original) + if capturedArgs != original { + t.Errorf("non-task tool args should not be modified, got %q", capturedArgs) + } +} + +func TestTaskContextEnrichMiddleware_NilWhenDisabled(t *testing.T) { + mw := newTaskContextEnrichMiddleware("test", nil, -1, "") + if mw != nil { + t.Error("middleware should be nil when disabled") + } +} diff --git a/internal/multiagent/system_message_normalizer_middleware.go b/internal/multiagent/system_message_normalizer_middleware.go new file mode 100644 index 00000000..6739d202 --- /dev/null +++ b/internal/multiagent/system_message_normalizer_middleware.go @@ -0,0 +1,86 @@ +package multiagent + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// systemMessageNormalizerMiddleware merges duplicate role=system messages into a single +// leading system message before summarization and each ChatModel call. +type systemMessageNormalizerMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + phase string +} + +func newSystemMessageNormalizerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware { + return &systemMessageNormalizerMiddleware{logger: logger, phase: phase} +} + +func (m *systemMessageNormalizerMiddleware) 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 + } + before := countADKSystemMessages(state.Messages) + if before <= 1 { + return ctx, state, nil + } + normalized := normalizeSingleLeadingSystemMessage(state.Messages, "") + if len(normalized) == len(state.Messages) && countADKSystemMessages(normalized) >= before { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Info("eino system messages merged", + zap.String("phase", m.phase), + zap.Int("system_before", before), + zap.Int("system_after", countADKSystemMessages(normalized)), + zap.Int("messages_before", len(state.Messages)), + zap.Int("messages_after", len(normalized)), + ) + } + out := *state + out.Messages = normalized + return ctx, &out, nil +} + +func countADKSystemMessages(msgs []adk.Message) int { + n := 0 + for _, msg := range msgs { + if msg != nil && msg.Role == schema.System { + n++ + } + } + return n +} + +// stripADKSystemMessages removes all system messages. Use before runner.Run restart when +// genModelInput will prepend a fresh Instruction. +func stripADKSystemMessages(msgs []adk.Message) []adk.Message { + if len(msgs) == 0 { + return msgs + } + out := make([]adk.Message, 0, len(msgs)) + for _, msg := range msgs { + if msg == nil || msg.Role == schema.System { + continue + } + out = append(out, msg) + } + return out +} + +// mergeCollectedSystemMessages collapses multiple system messages into one (or none). +func mergeCollectedSystemMessages(systemMsgs []adk.Message) []adk.Message { + if len(systemMsgs) == 0 { + return nil + } + return normalizeSingleLeadingSystemMessage(systemMsgs, "") +} diff --git a/internal/multiagent/system_message_normalizer_middleware_test.go b/internal/multiagent/system_message_normalizer_middleware_test.go new file mode 100644 index 00000000..eaf8219e --- /dev/null +++ b/internal/multiagent/system_message_normalizer_middleware_test.go @@ -0,0 +1,75 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestStripADKSystemMessages(t *testing.T) { + in := []adk.Message{ + schema.SystemMessage("a"), + schema.UserMessage("u"), + schema.SystemMessage("b"), + schema.AssistantMessage("x", nil), + } + out := stripADKSystemMessages(in) + if len(out) != 2 { + t.Fatalf("got %d messages, want 2", len(out)) + } + if out[0].Role != schema.User || out[1].Role != schema.Assistant { + t.Fatalf("unexpected roles: %s, %s", out[0].Role, out[1].Role) + } +} + +func TestEinoMessagesForRunRestart_StripsSystemFromTrace(t *testing.T) { + holder := newModelFacingTraceHolder() + holder.storeFromState(&adk.ChatModelAgentState{Messages: []adk.Message{ + schema.SystemMessage("sys-1"), + schema.SystemMessage("sys-2"), + schema.UserMessage("task"), + }}) + msgs, src := einoMessagesForRunRestart(&einoADKRunLoopArgs{ModelFacingTrace: holder}, nil, nil, 0) + if src != einoRestartContextModelTrace { + t.Fatalf("source: got %q want model_trace", src) + } + if len(msgs) != 1 || msgs[0].Role != schema.User { + t.Fatalf("expected user-only restart msgs, got %+v", msgs) + } +} + +func TestSystemMessageNormalizerMiddleware_MergesDuplicates(t *testing.T) { + mw := newSystemMessageNormalizerMiddleware(nil, "test") + state := &adk.ChatModelAgentState{Messages: []adk.Message{ + schema.SystemMessage("a"), + schema.SystemMessage("b"), + schema.UserMessage("u"), + }} + _, out, err := mw.(*systemMessageNormalizerMiddleware).BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatal(err) + } + if countADKSystemMessages(out.Messages) != 1 { + t.Fatalf("want 1 system, got %d", countADKSystemMessages(out.Messages)) + } + if out.Messages[0].Content != "a\n\nb" { + t.Fatalf("merged content: %q", out.Messages[0].Content) + } +} + +func TestSystemMessageNormalizerMiddleware_NoOpSingleSystem(t *testing.T) { + mw := newSystemMessageNormalizerMiddleware(nil, "test") + state := &adk.ChatModelAgentState{Messages: []adk.Message{ + schema.SystemMessage("only"), + schema.UserMessage("u"), + }} + _, out, err := mw.(*systemMessageNormalizerMiddleware).BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatal(err) + } + if out != state { + t.Fatalf("expected same state pointer for no-op") + } +} diff --git a/internal/multiagent/tool_always_visible.go b/internal/multiagent/tool_always_visible.go new file mode 100644 index 00000000..151cccc2 --- /dev/null +++ b/internal/multiagent/tool_always_visible.go @@ -0,0 +1,72 @@ +package multiagent + +import ( + "strings" +) + +// expandAlwaysVisibleNameSet 将配置中的常驻工具名展开为可匹配运行时工具名的集合。 +// 支持:内置短名 read_file;外部 mcp::tool;运行时 mcp__tool(OpenAI/Eino 命名)。 +func expandAlwaysVisibleNameSet(names []string) map[string]struct{} { + set := make(map[string]struct{}, len(names)*3) + add := func(name string) { + n := strings.TrimSpace(strings.ToLower(name)) + if n == "" { + return + } + set[n] = struct{}{} + } + for _, raw := range names { + n := strings.TrimSpace(strings.ToLower(raw)) + if n == "" { + continue + } + add(n) + if mcp, tool, ok := strings.Cut(n, "::"); ok && mcp != "" && tool != "" { + // 外部工具用 mcp::tool 配置时只展开运行时 mcp__tool,避免短名误伤其它 MCP 同名工具。 + add(mcp + "__" + tool) + continue + } + if idx := strings.LastIndex(n, "__"); idx > 0 { + mcp, tool := n[:idx], n[idx+2:] + if mcp != "" && tool != "" { + add(mcp + "::" + tool) + } + continue + } + } + return set +} + +// toolMatchesAlwaysVisible 判断运行时工具名是否命中常驻白名单(含别名)。 +func toolMatchesAlwaysVisible(runtimeName string, nameSet map[string]struct{}) bool { + if len(nameSet) == 0 { + return false + } + name := strings.TrimSpace(strings.ToLower(runtimeName)) + if name == "" { + return false + } + if _, ok := nameSet[name]; ok { + return true + } + if mcp, tool, ok := strings.Cut(name, "::"); ok && mcp != "" && tool != "" { + if _, ok := nameSet[mcp+"__"+tool]; ok { + return true + } + if _, ok := nameSet[tool]; ok { + return true + } + } + if idx := strings.LastIndex(name, "__"); idx > 0 { + mcp, tool := name[:idx], name[idx+2:] + if mcp != "" && tool != "" { + if _, ok := nameSet[mcp+"::"+tool]; ok { + return true + } + if _, ok := nameSet[tool]; ok { + return true + } + } + } + return false +} diff --git a/internal/multiagent/tool_always_visible_test.go b/internal/multiagent/tool_always_visible_test.go new file mode 100644 index 00000000..00c9eaa0 --- /dev/null +++ b/internal/multiagent/tool_always_visible_test.go @@ -0,0 +1,32 @@ +package multiagent + +import "testing" + +func TestToolMatchesAlwaysVisible_ExternalAliases(t *testing.T) { + t.Parallel() + set := expandAlwaysVisibleNameSet([]string{"zhidemai::discount_search", "read_file"}) + + cases := []struct { + runtime string + want bool + }{ + {"zhidemai__discount_search", true}, + {"zhidemai::discount_search", true}, + {"read_file", true}, + {"zhidemai__product_search_pro", false}, + {"github__discount_search", false}, + } + for _, tc := range cases { + if got := toolMatchesAlwaysVisible(tc.runtime, set); got != tc.want { + t.Fatalf("toolMatchesAlwaysVisible(%q) = %v, want %v", tc.runtime, got, tc.want) + } + } +} + +func TestExpandAlwaysVisibleNameSet_LegacyShortName(t *testing.T) { + t.Parallel() + set := expandAlwaysVisibleNameSet([]string{"discount_search"}) + if !toolMatchesAlwaysVisible("zhidemai__discount_search", set) { + t.Fatal("legacy short name should match external runtime tool") + } +} diff --git a/internal/multiagent/tool_call_arguments_sanitizer_middleware.go b/internal/multiagent/tool_call_arguments_sanitizer_middleware.go new file mode 100644 index 00000000..06e4b18f --- /dev/null +++ b/internal/multiagent/tool_call_arguments_sanitizer_middleware.go @@ -0,0 +1,103 @@ +package multiagent + +import ( + "context" + "encoding/json" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +const repairedMalformedToolArguments = `{}` + +// toolCallArgumentsSanitizerMiddleware guarantees that every historical +// tool_calls[].function.arguments value sent to an OpenAI-compatible provider is +// a syntactically valid JSON object. Some providers reject the entire request +// with HTTP 400 when a model previously emitted truncated arguments. +// +// The original malformed payload is intentionally not copied into model-facing +// history: it may contain secrets and can itself be large enough to trigger the +// same failure again. The paired tool result already records the execution error +// and gives the model enough information to recover. +type toolCallArgumentsSanitizerMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + phase string +} + +func newToolCallArgumentsSanitizerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware { + return &toolCallArgumentsSanitizerMiddleware{logger: logger, phase: phase} +} + +func (m *toolCallArgumentsSanitizerMiddleware) 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 + } + + out, repaired := sanitizeMalformedToolCallArguments(state.Messages) + if repaired == 0 { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Warn("eino malformed tool-call arguments repaired before model call", + zap.String("phase", m.phase), + zap.Int("repaired_calls", repaired), + ) + } + ns := *state + ns.Messages = out + return ctx, &ns, nil +} + +func sanitizeMalformedToolCallArguments(messages []adk.Message) ([]adk.Message, int) { + var out []adk.Message + repaired := 0 + for i, msg := range messages { + if msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + continue + } + calls := append([]schema.ToolCall(nil), msg.ToolCalls...) + changed := false + for j := range calls { + if validToolArgumentsJSONObject(calls[j].Function.Arguments) { + continue + } + calls[j].Function.Arguments = repairedMalformedToolArguments + changed = true + repaired++ + } + if !changed { + continue + } + if out == nil { + out = append([]adk.Message(nil), messages...) + } + cloned := *msg + cloned.ToolCalls = calls + out[i] = &cloned + } + if out == nil { + return messages, 0 + } + return out, repaired +} + +func validToolArgumentsJSONObject(arguments string) bool { + arguments = strings.TrimSpace(arguments) + if arguments == "" { + return false + } + var value any + if err := json.Unmarshal([]byte(arguments), &value); err != nil { + return false + } + _, ok := value.(map[string]any) + return ok +} diff --git a/internal/multiagent/tool_call_arguments_sanitizer_middleware_test.go b/internal/multiagent/tool_call_arguments_sanitizer_middleware_test.go new file mode 100644 index 00000000..361d78ed --- /dev/null +++ b/internal/multiagent/tool_call_arguments_sanitizer_middleware_test.go @@ -0,0 +1,65 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestToolCallArgumentsSanitizerRepairsOnlyMalformedObjects(t *testing.T) { + valid := assistantToolCallsMsg("", "valid") + valid.ToolCalls[0].Function.Arguments = `{"command":"echo ok"}` + malformed := assistantToolCallsMsg("", "broken", "array") + malformed.ToolCalls[0].Function.Arguments = `{"command":"unterminated` + malformed.ToolCalls[1].Function.Arguments = `[]` + messages := []adk.Message{valid, malformed, schema.ToolMessage("failed", "broken")} + + out, repaired := sanitizeMalformedToolCallArguments(messages) + if repaired != 2 { + t.Fatalf("repaired=%d, want 2", repaired) + } + if out[0].ToolCalls[0].Function.Arguments != `{"command":"echo ok"}` { + t.Fatalf("valid arguments changed: %q", out[0].ToolCalls[0].Function.Arguments) + } + for _, tc := range out[1].ToolCalls { + if tc.Function.Arguments != repairedMalformedToolArguments { + t.Fatalf("malformed arguments not repaired: %q", tc.Function.Arguments) + } + } + if malformed.ToolCalls[0].Function.Arguments == repairedMalformedToolArguments { + t.Fatal("input message was mutated") + } +} + +func TestToolCallArgumentsSanitizerMiddlewareRewritesState(t *testing.T) { + msg := assistantToolCallsMsg("", "broken") + msg.ToolCalls[0].Function.Arguments = "" + mw := newToolCallArgumentsSanitizerMiddleware(nil, "test").(*toolCallArgumentsSanitizerMiddleware) + _, state, err := mw.BeforeModelRewriteState(context.Background(), &adk.ChatModelAgentState{ + Messages: []adk.Message{msg}, + }, &adk.ModelContext{}) + if err != nil { + t.Fatal(err) + } + if got := state.Messages[0].ToolCalls[0].Function.Arguments; got != `{}` { + t.Fatalf("arguments=%q, want {}", got) + } +} + +func TestValidToolArgumentsJSONObject(t *testing.T) { + cases := map[string]bool{ + `{}`: true, + `{"x":1}`: true, + `null`: false, + `[]`: false, + `{"x":`: false, + ``: false, + } + for input, want := range cases { + if got := validToolArgumentsJSONObject(input); got != want { + t.Errorf("validToolArgumentsJSONObject(%q)=%v, want %v", input, got, want) + } + } +} diff --git a/internal/multiagent/tool_error_middleware.go b/internal/multiagent/tool_error_middleware.go new file mode 100644 index 00000000..899faeb7 --- /dev/null +++ b/internal/multiagent/tool_error_middleware.go @@ -0,0 +1,148 @@ +package multiagent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// softRecoveryToolCallMiddleware returns an InvokableToolMiddleware that catches +// specific recoverable errors from tool execution (JSON parse errors, tool-not-found, +// etc.) and converts them into soft errors: nil error + descriptive error content +// returned to the LLM. This allows the model to self-correct within the same +// iteration rather than crashing the entire graph and requiring a full replay. +// +// Without Invokable (+ Streamable where applicable) registration, a JSON parse failure +// in InvokableRun / StreamableRun propagates as a hard error through the Eino ToolsNode +// → [NodeRunError] → ev.Err, which +// either triggers the full-replay retry loop (expensive) or terminates the run +// entirely once retries are exhausted. With it, the LLM simply sees an error message +// in the tool result and can adjust its next tool call accordingly. +func softRecoveryToolCallMiddleware() compose.InvokableToolMiddleware { + return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + output, err := next(ctx, input) + if err == nil { + return output, nil + } + if !isSoftRecoverableToolError(err) { + return output, err + } + // Convert the hard error into a soft error: the LLM will see this + // message as the tool's output and can self-correct. + msg := buildSoftRecoveryMessage(input.Name, input.Arguments, err) + return &compose.ToolOutput{Result: msg}, nil + } + } +} + +// softRecoveryStreamableToolCallMiddleware mirrors softRecoveryToolCallMiddleware for +// tools that implement StreamableTool only (e.g. Eino ADK filesystem execute). +// Eino applies Invokable vs Streamable middleware to disjoint code paths in ToolsNode; +// registering only Invokable leaves streaming tools uncovered — empty/malformed JSON +// then fails inside [LocalStreamFunc] before the inner endpoint runs. +func softRecoveryStreamableToolCallMiddleware() compose.StreamableToolMiddleware { + return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + out, err := next(ctx, input) + if err == nil { + return out, nil + } + if !isSoftRecoverableToolError(err) { + return out, err + } + toolName := "" + args := "" + if input != nil { + toolName = input.Name + args = input.Arguments + } + msg := buildSoftRecoveryMessage(toolName, args, err) + return &compose.StreamToolOutput{ + Result: schema.StreamReaderFromArray([]string{msg}), + }, nil + } + } +} + +// softRecoveryToolMiddleware returns a ToolMiddleware with both Invokable and Streamable +// soft recovery (same semantics as hitlToolCallMiddleware bundling). +func softRecoveryToolMiddleware() compose.ToolMiddleware { + return compose.ToolMiddleware{ + Invokable: softRecoveryToolCallMiddleware(), + Streamable: softRecoveryStreamableToolCallMiddleware(), + } +} + +// isSoftRecoverableToolError determines whether a tool execution error should be +// silently converted to a tool-result message rather than crashing the graph. +// +// Design: default-soft (blacklist). Almost every tool execution error should be +// fed back to the LLM so it can self-correct or choose an alternative tool. +// Only a small set of "truly fatal" conditions (user cancellation) should +// propagate as hard errors that terminate the orchestration graph. +// This avoids the fragile whitelist approach where every new error pattern +// would need to be explicitly enumerated. +func isSoftRecoverableToolError(err error) bool { + if err == nil { + return false + } + + // 用户主动取消 — 唯一应当终止编排的情况,不应重试。 + if errors.Is(err, context.Canceled) { + return false + } + + // 其他所有工具执行错误(超时、命令不存在、JSON 解析失败、工具未找到、 + // 权限不足、网络不可达……)一律转为 soft error,让 LLM 看到错误信息 + // 后自行决策:换工具、调整参数、或向用户说明。 + return true +} + +// buildSoftRecoveryMessage creates a bilingual error message that the LLM can act on. +func buildSoftRecoveryMessage(toolName, arguments string, err error) string { + // Truncate arguments preview to avoid flooding the context. + argPreview := arguments + if len(argPreview) > 300 { + argPreview = argPreview[:300] + "... (truncated)" + } + + // Try to determine if it's specifically a JSON parse error for a friendlier message. + errStr := err.Error() + var jsonErr *json.SyntaxError + isJSONErr := strings.Contains(strings.ToLower(errStr), "json") || + strings.Contains(strings.ToLower(errStr), "unmarshal") + _ = jsonErr // suppress unused + + if isJSONErr { + return fmt.Sprintf( + "[Tool Error] The arguments for tool '%s' are not valid JSON and could not be parsed.\n"+ + "Error: %s\n"+ + "Arguments received: %s\n\n"+ + "Please fix the JSON (ensure double-quoted keys, matched braces/brackets, no trailing commas, "+ + "no truncation) and call the tool again.\n\n"+ + "[工具错误] 工具 '%s' 的参数不是合法 JSON,无法解析。\n"+ + "错误:%s\n"+ + "收到的参数:%s\n\n"+ + "请修正 JSON(确保双引号键名、括号配对、无尾部逗号、无截断),然后重新调用工具。", + toolName, errStr, argPreview, + toolName, errStr, argPreview, + ) + } + + return fmt.Sprintf( + "[Tool Error] Tool '%s' execution failed: %s\n"+ + "Arguments: %s\n\n"+ + "Please review the available tools and their expected arguments, then retry.\n\n"+ + "[工具错误] 工具 '%s' 执行失败:%s\n"+ + "参数:%s\n\n"+ + "请检查可用工具及其参数要求,然后重试。", + toolName, errStr, argPreview, + toolName, errStr, argPreview, + ) +} diff --git a/internal/multiagent/tool_error_middleware_test.go b/internal/multiagent/tool_error_middleware_test.go new file mode 100644 index 00000000..37e4fd70 --- /dev/null +++ b/internal/multiagent/tool_error_middleware_test.go @@ -0,0 +1,207 @@ +package multiagent + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/cloudwego/eino/compose" +) + +func TestIsSoftRecoverableToolError(t *testing.T) { + tests := []struct { + name string + err error + expected bool + }{ + { + name: "nil error", + err: nil, + expected: false, + }, + { + name: "unexpected end of JSON input", + err: errors.New("unexpected end of JSON input"), + expected: true, + }, + { + name: "failed to unmarshal task tool input json", + err: errors.New("failed to unmarshal task tool input json: unexpected end of JSON input"), + expected: true, + }, + { + name: "invalid tool arguments JSON", + err: errors.New("invalid tool arguments JSON: unexpected end of JSON input"), + expected: true, + }, + { + name: "json invalid character", + err: errors.New(`invalid character '}' looking for beginning of value in JSON`), + expected: true, + }, + { + name: "subagent type not found", + err: errors.New("subagent type recon_agent not found"), + expected: true, + }, + { + name: "tool not found", + err: errors.New("tool nmap_scan not found in toolsNode indexes"), + expected: true, + }, + { + name: "unrelated network error", + err: errors.New("connection refused"), + expected: true, // default-soft: non-cancel errors are recoverable + }, + { + name: "tool binary not installed", + err: errors.New("[LocalFunc] failed to invoke tool, toolName=grep, err=ripgrep (rg) is not installed or not in PATH"), + expected: true, + }, + { + name: "context cancelled", + err: context.Canceled, + expected: false, + }, + { + name: "real json unmarshal error", + err: func() error { + var v map[string]interface{} + return json.Unmarshal([]byte(`{"key": `), &v) + }(), + expected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := isSoftRecoverableToolError(tt.err) + if got != tt.expected { + t.Errorf("isSoftRecoverableToolError(%v) = %v, want %v", tt.err, got, tt.expected) + } + }) + } +} + +func TestSoftRecoveryToolCallMiddleware_PassesThrough(t *testing.T) { + mw := softRecoveryToolCallMiddleware() + called := false + next := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + called = true + return &compose.ToolOutput{Result: "success"}, nil + } + wrapped := mw(next) + out, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "test_tool", + Arguments: `{"key": "value"}`, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !called { + t.Fatal("next endpoint was not called") + } + if out.Result != "success" { + t.Fatalf("expected 'success', got %q", out.Result) + } +} + +func TestSoftRecoveryStreamableToolCallMiddleware_LocalStreamFuncJSONError(t *testing.T) { + mw := softRecoveryStreamableToolCallMiddleware() + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return nil, errors.New(`[LocalStreamFunc] failed to unmarshal arguments in json, toolName=execute, err="Syntax error no sources available, the input json is empty`) + } + wrapped := mw(next) + out, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "execute", + Arguments: "", + }) + if err != nil { + t.Fatalf("expected nil error (soft recovery), got: %v", err) + } + if out == nil || out.Result == nil { + t.Fatal("expected stream result") + } + var sb strings.Builder + for { + chunk, rerr := out.Result.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + sb.WriteString(chunk) + } + text := sb.String() + if !containsAll(text, "[Tool Error]", "execute", "JSON") { + t.Fatalf("recovery message missing expected content: %s", text) + } +} + +func TestSoftRecoveryToolCallMiddleware_ConvertsJSONError(t *testing.T) { + mw := softRecoveryToolCallMiddleware() + next := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return nil, errors.New("failed to unmarshal task tool input json: unexpected end of JSON input") + } + wrapped := mw(next) + out, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "task", + Arguments: `{"subagent_type": "recon`, + }) + if err != nil { + t.Fatalf("expected nil error (soft recovery), got: %v", err) + } + if out == nil || out.Result == "" { + t.Fatal("expected non-empty recovery message") + } + if !containsAll(out.Result, "[Tool Error]", "task", "JSON") { + t.Fatalf("recovery message missing expected content: %s", out.Result) + } +} + +func TestSoftRecoveryToolCallMiddleware_PropagatesNonRecoverable(t *testing.T) { + mw := softRecoveryToolCallMiddleware() + origErr := errors.New("connection timeout to remote server") + next := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return nil, origErr + } + wrapped := mw(next) + out, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "test_tool", + Arguments: `{}`, + }) + // Default-soft: non-cancel errors are converted to tool-result messages. + if err != nil { + t.Fatalf("expected nil error (soft recovery), got: %v", err) + } + if out == nil || out.Result == "" { + t.Fatal("expected non-empty recovery message") + } +} + +func containsAll(s string, subs ...string) bool { + for _, sub := range subs { + if !contains(s, sub) { + return false + } + } + return true +} + +func contains(s, sub string) bool { + return len(s) >= len(sub) && searchString(s, sub) +} + +func searchString(s, sub string) bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/internal/multiagent/tool_pair_reconciler_middleware.go b/internal/multiagent/tool_pair_reconciler_middleware.go new file mode 100644 index 00000000..fcd157b8 --- /dev/null +++ b/internal/multiagent/tool_pair_reconciler_middleware.go @@ -0,0 +1,156 @@ +package multiagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +const patchedMissingToolResult = "[Tool execution result was lost or interrupted; continue without relying on this call.]" + +// toolPairReconcilerMiddleware is the final structural guard before a model call. +// It makes every assistant tool-call batch immediately followed by exactly one tool +// result per call ID, and drops tool messages that cannot belong to that batch. +// +// This intentionally runs after summarization/reduction/budget middleware: those +// middlewares rewrite history and can otherwise re-introduce a partial tool round +// after the ordinary patchtoolcalls middleware has already run. +type toolPairReconcilerMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + phase string +} + +func newToolPairReconcilerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware { + return &toolPairReconcilerMiddleware{logger: logger, phase: phase} +} + +func (m *toolPairReconcilerMiddleware) 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 + } + + usedIDs := make(map[string]struct{}, 16) + changed := false + patched := 0 + dropped := 0 + out := make([]adk.Message, 0, len(state.Messages)) + + for i := 0; i < len(state.Messages); { + msg := state.Messages[i] + if msg == nil { + changed = true + i++ + continue + } + if msg.Role == schema.Tool { + // Valid tool results are consumed with their immediately preceding assistant. + changed = true + dropped++ + i++ + continue + } + if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + out = append(out, msg) + i++ + continue + } + + assistant := msg + calls := append([]schema.ToolCall(nil), msg.ToolCalls...) + expected := make(map[string]schema.ToolCall, len(calls)) + idsChanged := false + for callIndex := range calls { + id := calls[callIndex].ID + _, duplicate := usedIDs[id] + if id == "" || duplicate { + base := fmt.Sprintf("patched_tool_call_%d_%d", i, callIndex) + id = base + for suffix := 1; ; suffix++ { + if _, exists := usedIDs[id]; !exists { + break + } + id = fmt.Sprintf("%s_%d", base, suffix) + } + calls[callIndex].ID = id + idsChanged = true + changed = true + } + usedIDs[id] = struct{}{} + expected[id] = calls[callIndex] + } + if idsChanged { + cloned := *msg + cloned.ToolCalls = calls + assistant = &cloned + } + out = append(out, assistant) + + results := make(map[string]adk.Message, len(calls)) + j := i + 1 + for j < len(state.Messages) { + toolMsg := state.Messages[j] + if toolMsg == nil { + changed = true + j++ + continue + } + if toolMsg.Role != schema.Tool { + break + } + id := toolMsg.ToolCallID + if _, wanted := expected[id]; !wanted { + changed = true + dropped++ + j++ + continue + } + if _, duplicate := results[id]; duplicate { + changed = true + dropped++ + j++ + continue + } + results[id] = toolMsg + j++ + } + for _, tc := range calls { + if result, ok := results[tc.ID]; ok { + out = append(out, result) + continue + } + out = append(out, schema.ToolMessage( + patchedMissingToolResult, + tc.ID, + schema.WithToolName(tc.Function.Name), + )) + changed = true + patched++ + } + i = j + } + + if !changed { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Warn("eino tool-call/result pairs reconciled before model call", + zap.String("phase", m.phase), + zap.Int("patched_results", patched), + zap.Int("dropped_results", dropped), + zap.Int("messages_before", len(state.Messages)), + zap.Int("messages_after", len(out)), + ) + } + ns := *state + ns.Messages = out + return ctx, &ns, nil +} diff --git a/internal/multiagent/tool_pair_reconciler_middleware_test.go b/internal/multiagent/tool_pair_reconciler_middleware_test.go new file mode 100644 index 00000000..fac59ab0 --- /dev/null +++ b/internal/multiagent/tool_pair_reconciler_middleware_test.go @@ -0,0 +1,141 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func runToolPairReconciler(t *testing.T, msgs []adk.Message) []adk.Message { + t.Helper() + mw := newToolPairReconcilerMiddleware(nil, "test").(*toolPairReconcilerMiddleware) + _, out, err := mw.BeforeModelRewriteState( + context.Background(), + &adk.ChatModelAgentState{Messages: msgs}, + &adk.ModelContext{}, + ) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + return out.Messages +} + +func assertCompleteImmediateToolPairs(t *testing.T, msgs []adk.Message) { + t.Helper() + for i, msg := range msgs { + if msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + continue + } + want := make(map[string]struct{}, len(msg.ToolCalls)) + for _, tc := range msg.ToolCalls { + if tc.ID == "" { + t.Fatal("empty tool call id remained") + } + if _, duplicate := want[tc.ID]; duplicate { + t.Fatalf("duplicate tool call id remained: %s", tc.ID) + } + want[tc.ID] = struct{}{} + } + seen := make(map[string]struct{}, len(want)) + for j := i + 1; j < len(msgs) && msgs[j] != nil && msgs[j].Role == schema.Tool; j++ { + id := msgs[j].ToolCallID + if _, ok := want[id]; !ok { + t.Fatalf("unexpected tool result %q after assistant %d", id, i) + } + if _, duplicate := seen[id]; duplicate { + t.Fatalf("duplicate tool result %q", id) + } + seen[id] = struct{}{} + } + if len(seen) != len(want) { + t.Fatalf("assistant %d: want %d results, got %d", i, len(want), len(seen)) + } + } +} + +func TestToolPairReconcilerPatchesPartialMultiToolBatch(t *testing.T) { + msgs := []adk.Message{ + schema.UserMessage("start"), + assistantToolCallsMsg("", "c1", "c2"), + schema.ToolMessage("r1", "c1"), + schema.UserMessage("continue"), + } + out := runToolPairReconciler(t, msgs) + assertCompleteImmediateToolPairs(t, out) + if len(out) != 5 || out[3].Role != schema.Tool || out[3].ToolCallID != "c2" { + t.Fatalf("missing c2 result was not inserted in place: %+v", out) + } + if out[3].Content != patchedMissingToolResult { + t.Fatalf("unexpected patched content: %q", out[3].Content) + } +} + +func TestToolPairReconcilerDropsMisplacedDuplicateAndOrphanResults(t *testing.T) { + msgs := []adk.Message{ + schema.ToolMessage("old", "orphan"), + assistantToolCallsMsg("", "c1"), + schema.ToolMessage("first", "c1"), + schema.ToolMessage("duplicate", "c1"), + schema.ToolMessage("wrong", "other"), + schema.UserMessage("next"), + schema.ToolMessage("late", "c1"), + } + out := runToolPairReconciler(t, msgs) + assertCompleteImmediateToolPairs(t, out) + toolCount := 0 + for _, msg := range out { + if msg.Role == schema.Tool { + toolCount++ + if msg.Content != "first" || msg.ToolCallID != "c1" { + t.Fatalf("unexpected retained tool result: %+v", msg) + } + } + } + if toolCount != 1 { + t.Fatalf("want one retained tool result, got %d", toolCount) + } +} + +func TestToolPairReconcilerRepairsEmptyAndRepeatedCallIDs(t *testing.T) { + msgs := []adk.Message{ + assistantToolCallsMsg("", "", "same"), + schema.ToolMessage("same-1", "same"), + assistantToolCallsMsg("", "same"), + schema.ToolMessage("same-2", "same"), + } + out := runToolPairReconciler(t, msgs) + assertCompleteImmediateToolPairs(t, out) + all := make(map[string]struct{}) + for _, msg := range out { + if msg.Role != schema.Assistant { + continue + } + for _, tc := range msg.ToolCalls { + if _, duplicate := all[tc.ID]; duplicate { + t.Fatalf("global duplicate tool call id remained: %s", tc.ID) + } + all[tc.ID] = struct{}{} + } + } +} + +func TestToolPairReconcilerNoOpForValidHistory(t *testing.T) { + msgs := []adk.Message{ + schema.UserMessage("start"), + assistantToolCallsMsg("", "c1", "c2"), + schema.ToolMessage("r1", "c1"), + schema.ToolMessage("r2", "c2"), + schema.AssistantMessage("done", nil), + } + mw := newToolPairReconcilerMiddleware(nil, "test").(*toolPairReconcilerMiddleware) + in := &adk.ChatModelAgentState{Messages: msgs} + _, out, err := mw.BeforeModelRewriteState(context.Background(), in, &adk.ModelContext{}) + if err != nil { + t.Fatal(err) + } + if out != in { + t.Fatal("valid history should use the no-op fast path") + } +} diff --git a/internal/multiagent/tool_search_result_sanitizer_middleware.go b/internal/multiagent/tool_search_result_sanitizer_middleware.go new file mode 100644 index 00000000..df9bb247 --- /dev/null +++ b/internal/multiagent/tool_search_result_sanitizer_middleware.go @@ -0,0 +1,76 @@ +package multiagent + +import ( + "context" + "encoding/json" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// toolSearchResultSanitizerMiddleware prevents malformed historical tool_search +// results (for example an HTML gateway error page) from crashing Eino's dynamic +// tool loader on every retry. Eino expects every tool_search result to be a JSON +// object containing selectedTools. +type toolSearchResultSanitizerMiddleware struct { + adk.BaseChatModelAgentMiddleware + logger *zap.Logger + phase string +} + +func newToolSearchResultSanitizerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware { + return &toolSearchResultSanitizerMiddleware{logger: logger, phase: phase} +} + +type toolSearchResultEnvelope struct { + SelectedTools []string `json:"selectedTools"` +} + +func validToolSearchResult(content string) bool { + var result toolSearchResultEnvelope + if err := json.Unmarshal([]byte(content), &result); err != nil { + return false + } + // Reject JSON values such as null. They unmarshal without an error but do not + // satisfy the object-shaped contract used by the toolsearch middleware. + return strings.HasPrefix(strings.TrimSpace(content), "{") +} + +func (m *toolSearchResultSanitizerMiddleware) BeforeModelRewriteState( + ctx context.Context, + state *adk.ChatModelAgentState, + _ *adk.ModelContext, +) (context.Context, *adk.ChatModelAgentState, error) { + if m == nil || state == nil || len(state.Messages) == 0 { + return ctx, state, nil + } + + var rewritten []adk.Message + repaired := 0 + for i, msg := range state.Messages { + if msg == nil || msg.Role != schema.Tool || !IsToolSearchTool(msg.ToolName) || validToolSearchResult(msg.Content) { + continue + } + if rewritten == nil { + rewritten = append([]adk.Message(nil), state.Messages...) + } + clone := *msg + clone.Content = `{"selectedTools":[],"_recovered":true,"reason":"invalid historical tool_search result"}` + rewritten[i] = &clone + repaired++ + } + + if repaired == 0 { + return ctx, state, nil + } + if m.logger != nil { + m.logger.Warn("invalid historical tool_search results repaired before model call", + zap.String("phase", m.phase), + zap.Int("repaired_count", repaired)) + } + ns := *state + ns.Messages = rewritten + return ctx, &ns, nil +} diff --git a/internal/multiagent/tool_search_result_sanitizer_middleware_test.go b/internal/multiagent/tool_search_result_sanitizer_middleware_test.go new file mode 100644 index 00000000..c5b026cc --- /dev/null +++ b/internal/multiagent/tool_search_result_sanitizer_middleware_test.go @@ -0,0 +1,56 @@ +package multiagent + +import ( + "context" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestToolSearchResultSanitizerRepairsMalformedHistory(t *testing.T) { + good := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: `{"selectedTools":["grep"]}`} + bad := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: "502 Bad Gateway"} + other := &schema.Message{Role: schema.Tool, ToolName: "grep", Content: "plain text is valid for other tools"} + state := &adk.ChatModelAgentState{Messages: []adk.Message{good, bad, other}} + + mw := newToolSearchResultSanitizerMiddleware(nil, "test") + _, got, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + if got.Messages[0] != good || got.Messages[0].Content != good.Content { + t.Fatal("valid tool_search result was unexpectedly changed") + } + if got.Messages[1] == bad || !validToolSearchResult(got.Messages[1].Content) { + t.Fatalf("malformed result was not safely replaced: %q", got.Messages[1].Content) + } + if got.Messages[2] != other { + t.Fatal("non-tool_search result was unexpectedly changed") + } + if bad.Content != "502 Bad Gateway" { + t.Fatal("middleware mutated the original message") + } +} + +func TestToolSearchResultSanitizerFastPath(t *testing.T) { + msg := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: `{"selectedTools":[]}`} + state := &adk.ChatModelAgentState{Messages: []adk.Message{msg}} + mw := newToolSearchResultSanitizerMiddleware(nil, "test") + + _, got, err := mw.BeforeModelRewriteState(context.Background(), state, nil) + if err != nil { + t.Fatalf("BeforeModelRewriteState: %v", err) + } + if got != state { + t.Fatal("valid history should use the allocation-free fast path") + } +} + +func TestValidToolSearchResultRejectsNonObjectJSON(t *testing.T) { + for _, content := range []string{"null", `[]`, `"text"`, `{"selectedTools":"grep"}`} { + if validToolSearchResult(content) { + t.Fatalf("expected invalid tool_search result: %s", content) + } + } +}