From 5720de0d27dc67638497af08ff10c676a5d358b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:36:34 +0800 Subject: [PATCH] Add files via upload --- internal/attackchain/builder.go | 952 +++++++++++ internal/attackchain/promote_project.go | 203 +++ internal/attackchain/truncate.go | 248 +++ internal/attackchain/truncate_test.go | 63 + internal/audit/conversation_create.go | 55 + internal/audit/meta.go | 9 + internal/audit/record.go | 29 + internal/audit/resource_availability.go | 86 + internal/audit/retention.go | 27 + internal/audit/sanitize.go | 58 + internal/audit/service.go | 177 ++ internal/audit/throttle.go | 55 + internal/audit/types.go | 16 + internal/einomcp/holder.go | 21 + internal/einomcp/mcp_tools.go | 214 +++ internal/einomcp/mcp_tools_test.go | 16 + internal/einomcp/tool_invoke_notify.go | 39 + internal/einoobserve/attach.go | 455 +++++ internal/einoobserve/attach_test.go | 49 + internal/einoobserve/otel.go | 111 ++ internal/logger/logger.go | 68 + internal/mcp/builtin/constants.go | 195 +++ internal/mcp/client_sdk.go | 475 ++++++ internal/mcp/connection_recovery.go | 192 +++ internal/mcp/connection_recovery_test.go | 215 +++ internal/mcp/execution_control_tools.go | 296 ++++ internal/mcp/execution_service.go | 625 +++++++ internal/mcp/execution_service_test.go | 41 + internal/mcp/external_manager.go | 1615 ++++++++++++++++++ internal/mcp/external_manager_async_test.go | 230 +++ internal/mcp/external_manager_test.go | 261 +++ internal/mcp/run_context.go | 147 ++ internal/mcp/server.go | 1704 +++++++++++++++++++ internal/mcp/server_authorization_test.go | 231 +++ internal/mcp/tool_result_guard.go | 65 + internal/mcp/tool_result_guard_test.go | 158 ++ internal/mcp/types.go | 338 ++++ internal/projectprompt/blackboard.go | 132 ++ internal/projectprompt/shell_tools.go | 11 + internal/reasoning/eino.go | 428 +++++ internal/reasoning/eino_test.go | 424 +++++ internal/tooloutput/spill.go | 292 ++++ internal/tooloutput/spill_test.go | 58 + internal/vision/client.go | 166 ++ internal/vision/client_test.go | 12 + internal/vision/path.go | 72 + internal/vision/path_test.go | 52 + internal/vision/preprocess.go | 212 +++ internal/vision/preprocess_test.go | 109 ++ internal/vision/tool.go | 125 ++ 50 files changed, 11832 insertions(+) create mode 100644 internal/attackchain/builder.go create mode 100644 internal/attackchain/promote_project.go create mode 100644 internal/attackchain/truncate.go create mode 100644 internal/attackchain/truncate_test.go create mode 100644 internal/audit/conversation_create.go create mode 100644 internal/audit/meta.go create mode 100644 internal/audit/record.go create mode 100644 internal/audit/resource_availability.go create mode 100644 internal/audit/retention.go create mode 100644 internal/audit/sanitize.go create mode 100644 internal/audit/service.go create mode 100644 internal/audit/throttle.go create mode 100644 internal/audit/types.go create mode 100644 internal/einomcp/holder.go create mode 100644 internal/einomcp/mcp_tools.go create mode 100644 internal/einomcp/mcp_tools_test.go create mode 100644 internal/einomcp/tool_invoke_notify.go create mode 100644 internal/einoobserve/attach.go create mode 100644 internal/einoobserve/attach_test.go create mode 100644 internal/einoobserve/otel.go create mode 100644 internal/logger/logger.go create mode 100644 internal/mcp/builtin/constants.go create mode 100644 internal/mcp/client_sdk.go create mode 100644 internal/mcp/connection_recovery.go create mode 100644 internal/mcp/connection_recovery_test.go create mode 100644 internal/mcp/execution_control_tools.go create mode 100644 internal/mcp/execution_service.go create mode 100644 internal/mcp/execution_service_test.go create mode 100644 internal/mcp/external_manager.go create mode 100644 internal/mcp/external_manager_async_test.go create mode 100644 internal/mcp/external_manager_test.go create mode 100644 internal/mcp/run_context.go create mode 100644 internal/mcp/server.go create mode 100644 internal/mcp/server_authorization_test.go create mode 100644 internal/mcp/tool_result_guard.go create mode 100644 internal/mcp/tool_result_guard_test.go create mode 100644 internal/mcp/types.go create mode 100644 internal/projectprompt/blackboard.go create mode 100644 internal/projectprompt/shell_tools.go create mode 100644 internal/reasoning/eino.go create mode 100644 internal/reasoning/eino_test.go create mode 100644 internal/tooloutput/spill.go create mode 100644 internal/tooloutput/spill_test.go create mode 100644 internal/vision/client.go create mode 100644 internal/vision/client_test.go create mode 100644 internal/vision/path.go create mode 100644 internal/vision/path_test.go create mode 100644 internal/vision/preprocess.go create mode 100644 internal/vision/preprocess_test.go create mode 100644 internal/vision/tool.go diff --git a/internal/attackchain/builder.go b/internal/attackchain/builder.go new file mode 100644 index 00000000..f257f5d9 --- /dev/null +++ b/internal/attackchain/builder.go @@ -0,0 +1,952 @@ +package attackchain + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/openai" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Builder 攻击链构建器 +type Builder struct { + db *database.DB + logger *zap.Logger + openAIClient *openai.Client + openAIConfig *config.OpenAIConfig + tokenCounter agent.TokenCounter + maxTokens int // 最大tokens限制，默认100000 +} + +// Node 攻击链节点（使用database包的类型） +type Node = database.AttackChainNode + +// Edge 攻击链边（使用database包的类型） +type Edge = database.AttackChainEdge + +// Chain 完整的攻击链 +type Chain struct { + Nodes []Node `json:"nodes"` + Edges []Edge `json:"edges"` +} + +// NewBuilder 创建新的攻击链构建器 +func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.Logger) *Builder { + transport := &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + } + httpClient := &http.Client{Timeout: 5 * time.Minute, Transport: transport} + + // 优先使用配置文件中的统一 Token 上限（config.yaml -> openai.max_total_tokens） + maxTokens := 0 + if openAIConfig != nil && openAIConfig.MaxTotalTokens > 0 { + maxTokens = openAIConfig.MaxTotalTokens + } else if openAIConfig != nil { + // 如果未显式配置 max_total_tokens，则根据模型设置一个合理的默认值 + model := strings.ToLower(openAIConfig.Model) + if strings.Contains(model, "gpt-4") { + maxTokens = 128000 // gpt-4通常支持128k + } else if strings.Contains(model, "gpt-3.5") { + maxTokens = 16000 // gpt-3.5-turbo通常支持16k + } else if strings.Contains(model, "deepseek") { + maxTokens = 131072 // deepseek-chat通常支持131k + } else { + maxTokens = 100000 // 兜底默认值 + } + } else { + // 没有 OpenAI 配置时使用兜底值，避免为 0 + maxTokens = 100000 + } + + return &Builder{ + db: db, + logger: logger, + openAIClient: openai.NewClient(openAIConfig, httpClient, logger), + openAIConfig: openAIConfig, + tokenCounter: agent.NewTikTokenCounter(), + maxTokens: maxTokens, + } +} + +// BuildChainFromConversation 从对话构建攻击链（单次 LLM 调用；输入为当前任务轮次的 last_react 轨迹，与继续对话续跑范围一致）。 +func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID string) (*Chain, error) { + b.logger.Info("开始构建攻击链（简化版本）", zap.String("conversationId", conversationID)) + + // 0. 首先检查是否有实际的工具执行记录 + messages, err := b.db.GetMessages(conversationID) + if err != nil { + return nil, fmt.Errorf("获取对话消息失败: %w", err) + } + + if len(messages) == 0 { + b.logger.Info("对话中没有数据", zap.String("conversationId", conversationID)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 检查是否有实际的工具执行：assistant 的 mcp_execution_ids，或过程详情中的 tool_call/tool_result + //（多代理下若 MCP 未返回 execution_id，IDs 可能为空，但工具已通过 Eino 执行并写入 process_details） + hasToolExecutions := false + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + if len(messages[i].MCPExecutionIDs) > 0 { + hasToolExecutions = true + break + } + } + } + if !hasToolExecutions { + if pdOK, err := b.db.ConversationHasToolProcessDetails(conversationID); err != nil { + b.logger.Warn("查询过程详情判定工具执行失败", zap.Error(err)) + } else if pdOK { + hasToolExecutions = true + } + } + + // 检查任务是否被取消（通过检查最后一条assistant消息内容或process_details） + taskCancelled := false + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + content := strings.ToLower(messages[i].Content) + if strings.Contains(content, "取消") || strings.Contains(content, "cancelled") { + taskCancelled = true + } + break + } + } + + // 如果任务被取消且没有实际工具执行，返回空攻击链 + if taskCancelled && !hasToolExecutions { + b.logger.Info("任务已取消且没有实际工具执行，返回空攻击链", + zap.String("conversationId", conversationID), + zap.Bool("taskCancelled", taskCancelled), + zap.Bool("hasToolExecutions", hasToolExecutions)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 如果没有实际工具执行，也返回空攻击链（避免AI编造） + if !hasToolExecutions { + b.logger.Info("没有实际工具执行记录，返回空攻击链", + zap.String("conversationId", conversationID)) + return &Chain{Nodes: []Node{}, Edges: []Edge{}}, nil + } + + // 1. 优先尝试从数据库获取保存的最后一轮ReAct输入和输出 + reactInputJSON, modelOutput, err := b.db.GetAgentTrace(conversationID) + if err != nil { + b.logger.Warn("获取保存的ReAct数据失败，将使用消息历史构建", zap.Error(err)) + // 继续使用原来的逻辑 + reactInputJSON = "" + modelOutput = "" + } + + // var userInput string + var reactInputFinal string + var dataSource string // 记录数据来源 + + // 优先使用落库的代理轨迹（与继续对话 loadHistoryFromAgentTrace 同源），并裁剪为「当前任务轮次」 + if reactInputJSON != "" { + trimmedJSON := agent.ExtractLastUserTurnTraceJSON(reactInputJSON) + hash := sha256.Sum256([]byte(trimmedJSON)) + reactInputHash := hex.EncodeToString(hash[:])[:16] + + var messageCount int + if msgs, parseErr := agent.ParseTraceMessages(trimmedJSON); parseErr == nil { + messageCount = len(msgs) + msgs = agent.MergeAssistantTraceOutput(msgs, modelOutput) + reactInputFinal = b.formatAgentTraceFromChatMessages(msgs) + } else { + b.logger.Warn("解析代理轨迹失败，回退原始 JSON 格式化", zap.Error(parseErr)) + reactInputFinal = b.formatAgentTraceInputFromJSON(trimmedJSON) + if strings.TrimSpace(modelOutput) != "" { + reactInputFinal += "\n\n## 助手结论（last_react_output）\n\n" + modelOutput + } + } + + dataSource = "last_user_turn_agent_trace" + b.logger.Info("使用当前任务轮次代理轨迹构建攻击链（与续跑上下文范围一致）", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("traceInputSizeBeforeTrim", len(reactInputJSON)), + zap.Int("traceInputSizeAfterTrim", len(trimmedJSON)), + zap.Int("messageCount", messageCount), + zap.String("reactInputHash", reactInputHash), + zap.Int("modelOutputSize", len(modelOutput))) + } else { + // 2. 如果没有保存的ReAct数据，从对话消息构建 + dataSource = "messages_table" + b.logger.Info("从消息历史构建ReAct数据", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("messageCount", len(messages))) + + // 提取用户输入（最后一条user消息） + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "user") { + // userInput = messages[i].Content + break + } + } + + // 提取最后一轮ReAct的输入（历史消息+当前用户输入） + reactInputFinal = b.buildAgentTraceInput(messages) + + // 提取大模型最后的输出（最后一条assistant消息） + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + modelOutput = messages[i].Content + break + } + } + } + + // 多代理：保存的轨迹列可能仅为首轮用户消息，不含工具轨迹；补充最后一轮助手的过程详情（与单代理完整轨迹对齐） + hasMCPOnAssistant := false + var lastAssistantID string + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "assistant") { + lastAssistantID = messages[i].ID + if len(messages[i].MCPExecutionIDs) > 0 { + hasMCPOnAssistant = true + } + break + } + } + if lastAssistantID != "" { + pdHasTools, _ := b.db.ConversationHasToolProcessDetails(conversationID) + if pdHasTools && !(hasMCPOnAssistant && reactInputContainsToolTrace(reactInputJSON)) { + detailsMap, err := b.db.GetProcessDetailsByConversation(conversationID) + if err != nil { + b.logger.Warn("加载过程详情用于攻击链失败", zap.Error(err)) + } else if dets := detailsMap[lastAssistantID]; len(dets) > 0 { + extra := b.formatProcessDetailsForAttackChain(dets) + if strings.TrimSpace(extra) != "" { + reactInputFinal = reactInputFinal + "\n\n## 执行过程与工具记录（含多代理编排与子任务）\n\n" + extra + b.logger.Info("攻击链输入已补充过程详情", + zap.String("conversationId", conversationID), + zap.String("messageId", lastAssistantID), + zap.Int("detailEvents", len(dets))) + } + } + } + } + + // 3. 按 token 预算压缩输入，再构建 prompt（避免超出模型上下文） + reactInputFinal, modelOutput, _ = b.fitAttackChainPayload(reactInputFinal, modelOutput) + + // 4. 构建 prompt 并单次调用大模型（助手结论已并入轨迹时不再重复传入） + promptAssistantOut := modelOutput + if reactInputJSON != "" { + promptAssistantOut = "" + } + prompt := b.buildSimplePrompt(reactInputFinal, promptAssistantOut) + // fmt.Println(prompt) + // 6. 调用AI生成攻击链（一次性，不做任何处理） + chainJSON, err := b.callAIForChainGeneration(ctx, prompt) + if err != nil { + return nil, fmt.Errorf("AI生成失败: %w", err) + } + + // 7. 解析JSON并生成节点/边ID（前端需要有效的ID） + chainData, err := b.parseChainJSON(chainJSON) + if err != nil { + // 如果解析失败，返回空链，让前端处理错误 + b.logger.Warn("解析攻击链JSON失败", zap.Error(err), zap.String("raw_json", chainJSON)) + return &Chain{ + Nodes: []Node{}, + Edges: []Edge{}, + }, nil + } + + b.logger.Info("攻击链构建完成", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("nodes", len(chainData.Nodes)), + zap.Int("edges", len(chainData.Edges))) + + // 保存到数据库（供后续加载使用） + if err := b.saveChain(conversationID, chainData.Nodes, chainData.Edges); err != nil { + b.logger.Warn("保存攻击链到数据库失败", zap.Error(err)) + // 即使保存失败，也返回数据给前端 + } + + // 直接返回，不做任何处理和校验 + return chainData, nil +} + +// reactInputContainsToolTrace 判断保存的 ReAct JSON 是否包含可解析的工具调用轨迹（单代理完整保存时为 true）。 +func reactInputContainsToolTrace(reactInputJSON string) bool { + s := strings.TrimSpace(reactInputJSON) + if s == "" { + return false + } + return strings.Contains(s, "tool_calls") || + strings.Contains(s, "tool_call_id") || + strings.Contains(s, `"role":"tool"`) || + strings.Contains(s, `"role": "tool"`) +} + +// formatProcessDetailsForAttackChain 将最后一轮助手的过程详情格式化为攻击链分析的输入（覆盖多代理下 last_react_input 不完整的情况）。 +func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessDetail) string { + if len(details) == 0 { + return "" + } + var sb strings.Builder + for _, d := range details { + // 目标：以主 agent（编排器）视角输出整轮迭代 + // - 保留：编排器工具调用/结果、对子代理的 task 调度、子代理最终回复（不含推理） + // - 丢弃：thinking/planning/progress 等噪声、子代理的工具细节与推理过程 + if d.EventType == "progress" || d.EventType == "thinking" || d.EventType == "reasoning_chain" || d.EventType == "planning" { + continue + } + + // 解析 data（JSON string），用于识别 einoRole / toolName 等 + var dataMap map[string]interface{} + if strings.TrimSpace(d.Data) != "" { + _ = json.Unmarshal([]byte(d.Data), &dataMap) + } + einoRole := "" + if v, ok := dataMap["einoRole"]; ok { + einoRole = strings.ToLower(strings.TrimSpace(fmt.Sprint(v))) + } + toolName := "" + if v, ok := dataMap["toolName"]; ok { + toolName = strings.TrimSpace(fmt.Sprint(v)) + } + + // 1) 编排器的工具调用/结果：保留（这是“主 agent 调了什么工具”） + if (d.EventType == "tool_call" || d.EventType == "tool_result" || d.EventType == "tool_calls_detected" || d.EventType == "iteration") && einoRole == "orchestrator" { + sb.WriteString("[") + sb.WriteString(d.EventType) + sb.WriteString("] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 2) 子代理调度：tool_call(toolName=="task") 代表编排器把子任务派发出去；保留（只需任务，不要子代理推理） + if d.EventType == "tool_call" && strings.EqualFold(toolName, "task") { + sb.WriteString("[dispatch_subagent_task] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 3) 子代理最终回复：保留（只保留最终输出，不保留分析过程） + if d.EventType == "eino_agent_reply" && einoRole == "sub" { + sb.WriteString("[subagent_final_reply] ") + sb.WriteString(strings.TrimSpace(d.Message)) + sb.WriteString("\n") + // data 里含 einoAgent 等元信息，保留有助于追踪“哪个子代理说的” + if strings.TrimSpace(d.Data) != "" { + sb.WriteString(d.Data) + sb.WriteString("\n") + } + sb.WriteString("\n") + continue + } + + // 其他事件默认丢弃，避免把子代理工具细节/推理塞进 prompt，偏离“主 agent 一轮迭代”的视角。 + } + return strings.TrimSpace(sb.String()) +} + +// buildAgentTraceInput 构建最后一轮 ReAct 的输入（从最后一条 user 消息起，不含更早轮次）。 +func (b *Builder) buildAgentTraceInput(messages []database.Message) string { + start := 0 + for i := len(messages) - 1; i >= 0; i-- { + if strings.EqualFold(messages[i].Role, "user") { + start = i + break + } + } + var builder strings.Builder + for _, msg := range messages[start:] { + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content)) + } + return builder.String() +} + +// extractUserInputFromReActInput 从保存的ReAct输入（JSON格式的messages数组）中提取最后一条用户输入 +// func (b *Builder) extractUserInputFromReActInput(reactInputJSON string) string { +// // reactInputJSON是JSON格式的ChatMessage数组，需要解析 +// var messages []map[string]interface{} +// if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil { +// b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err)) +// return "" +// } + +// // 从后往前查找最后一条user消息 +// for i := len(messages) - 1; i >= 0; i-- { +// if role, ok := messages[i]["role"].(string); ok && strings.EqualFold(role, "user") { +// if content, ok := messages[i]["content"].(string); ok { +// return content +// } +// } +// } + +// return "" +// } + +// formatAgentTraceInputFromJSON 将 JSON 轨迹转为可读文本（会先按当前任务轮次裁剪）。 +func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string { + trimmed := agent.ExtractLastUserTurnTraceJSON(reactInputJSON) + msgs, err := agent.ParseTraceMessages(trimmed) + if err != nil { + b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err)) + return trimmed + } + return b.formatAgentTraceFromChatMessages(msgs) +} + +// formatAgentTraceFromChatMessages 将代理消息带格式化为攻击链分析输入（与续跑轨迹字段一致）。 +func (b *Builder) formatAgentTraceFromChatMessages(msgs []agent.ChatMessage) string { + var builder strings.Builder + for _, msg := range msgs { + role := msg.Role + content := msg.Content + + if strings.EqualFold(role, "assistant") && len(msg.ToolCalls) > 0 { + if content != "" { + builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content)) + } + builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(msg.ToolCalls))) + for i, tc := range msg.ToolCalls { + args := "" + if tc.Function.Arguments != nil { + if b, err := json.Marshal(tc.Function.Arguments); err == nil { + args = string(b) + } + } + builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1)) + builder.WriteString(fmt.Sprintf(" ID: %s\n", tc.ID)) + builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", tc.Function.Name)) + builder.WriteString(fmt.Sprintf(" 参数: %s\n", args)) + } + builder.WriteString("\n") + continue + } + + if strings.EqualFold(role, "tool") { + if msg.ToolCallID != "" { + builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, msg.ToolCallID, content)) + } else { + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content)) + } + continue + } + + builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content)) + } + return builder.String() +} + +// buildSimplePrompt 构建简化的prompt +func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string { + return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据**当前任务轮次**的对话记录和工具执行结果，一次性输出攻击链 JSON（不要分多轮追问）。 + +## 输入范围（与「继续对话」续跑一致） +- 下方「ReAct 轨迹」仅包含**最后一次用户提问之后**的消息与工具结果（last_react 当前任务轮次），不含更早的用户提问轮次。 +- 「助手结论」为同轮任务的最终输出摘要（last_react_output）；节点须与轨迹中的实际工具执行一致，严禁编造。 + +## 核心目标 + +构建一个能够讲述完整攻击故事的攻击链让学习者能够： +1. 理解渗透测试的完整流程和思维逻辑（从目标识别到漏洞发现的每一步） +2. 学习如何从失败中获取线索并调整策略 +3. 掌握工具使用的实际效果和局限性 +4. 理解漏洞发现和利用的因果关系 + +**关键原则**：完整性优先。必须包含所有有意义的工具执行和关键步骤，不要为了控制节点数量而遗漏重要信息。 + +## 构建流程（按此顺序思考） + +### 第一步：理解上下文 +仔细分析ReAct输入中的工具调用序列和大模型输出，识别： +- 测试目标（IP、域名、URL等） +- 实际执行的工具和参数 +- 工具返回的关键信息（成功结果、错误信息、超时等） +- AI的分析和决策过程 + +### 第二步：提取关键节点 +从工具执行记录中提取有意义的节点，**确保不遗漏任何关键步骤**： +- **target节点**：每个独立的测试目标创建一个target节点 +- **action节点**：每个有意义的工具执行创建一个action节点（包括提供线索的失败、成功的信息收集、漏洞验证等） +- **vulnerability节点**：每个真实确认的漏洞创建一个vulnerability节点 +- **完整性检查**：对照ReAct输入中的工具调用序列，确保每个有意义的工具执行都被包含在攻击链中 + +### 第三步：构建逻辑关系（树状结构） +**重要：必须构建树状结构，而不是简单的线性链。** +按照因果关系连接节点，形成树状图（因为是单agent执行，所以可以不按照时间顺序）： +- **分支结构**：一个节点可以有多个后续节点（例如：端口扫描发现多个端口后，可以同时进行多个不同的测试） +- **汇聚结构**：多个节点可以指向同一个节点（例如：多个不同的测试都发现了同一个漏洞） +- 识别哪些action是基于前面action的结果而执行的 +- 识别哪些vulnerability是由哪些action发现的 +- 识别失败节点如何为后续成功提供线索 +- **避免线性链**：不要将所有节点连成一条线，应该根据实际的并行测试和分支探索构建树状结构 + +### 第四步：优化和精简 +- **完整性检查**：确保所有有意义的工具执行都被包含，不要遗漏关键步骤 +- **合并规则**：只合并真正相似或重复的action节点（如多次相同工具的相似调用） +- **删除规则**：只删除完全无价值的失败节点（完全无输出、纯系统错误、重复的相同失败） +- **重要提醒**：宁可保留更多节点，也不要遗漏关键步骤。攻击链必须完整展现渗透测试过程 +- 确保攻击链逻辑连贯，能够讲述完整故事 + +## 节点类型详解 + +### target（目标节点） +- **用途**：标识测试目标 +- **创建规则**：每个独立目标（不同IP/域名）创建一个target节点 +- **多目标处理**：不同目标的节点不相互连接，各自形成独立的子图 +- **metadata.target**：精确记录目标标识（IP地址、域名、URL等） + +### action（行动节点） +- **用途**：记录工具执行和AI分析结果 +- **标签规则**： + * 15-25个汉字，动宾结构 + * 成功节点：描述执行结果（如"扫描端口发现80/443/8080"、"目录扫描发现/admin路径"） + * 失败节点：描述失败原因（如"尝试SQL注入（被WAF拦截）"、"端口扫描超时（目标不可达）"） +- **ai_analysis要求**： + * 成功节点：总结工具执行的关键发现，说明这些发现的意义 + * 失败节点：必须说明失败原因、获得的线索、这些线索如何指引后续行动 + * 不超过150字，要具体、有信息量 +- **findings要求**： + * 提取工具返回结果中的关键信息点 + * 每个finding应该是独立的、有价值的信息片段 + * 成功节点：列出关键发现（如["80端口开放", "443端口开放", "HTTP服务为Apache 2.4"]） + * 失败节点：列出失败线索（如["WAF拦截", "返回403", "检测到Cloudflare"]） +- **status标记**： + * 成功节点：不设置或设为"success" + * 提供线索的失败节点：必须设为"failed_insight" +- **risk_score**：始终为0（action节点不评估风险） + +### vulnerability（漏洞节点） +- **用途**：记录真实确认的安全漏洞 +- **创建规则**： + * 必须是真实确认的漏洞，不是所有发现都是漏洞 + * 需要明确的漏洞证据（如SQL注入返回数据库错误、XSS成功执行等） +- **risk_score规则**： + * critical（90-100）：可导致系统完全沦陷（RCE、SQL注入导致数据泄露等） + * high（80-89）：可导致敏感信息泄露或权限提升 + * medium（60-79）：存在安全风险但影响有限 + * low（40-59）：轻微安全问题 +- **metadata要求**： + * vulnerability_type：漏洞类型（SQL注入、XSS、RCE等） + * description：详细描述漏洞位置、原理、影响 + * severity：critical/high/medium/low + * location：精确的漏洞位置（URL、参数、文件路径等） + +## 节点过滤和合并规则 + +### 必须保留的失败节点 +以下失败情况必须创建节点，因为它们提供了有价值的线索： +- 工具返回明确的错误信息（权限错误、连接拒绝、认证失败等） +- 超时或连接失败（可能表明防火墙、网络隔离等） +- WAF/防火墙拦截（返回403、406等，表明存在防护机制） +- 工具未安装或配置错误（但执行了调用） +- 目标不可达（DNS解析失败、网络不通等） + +### 应该删除的失败节点 +以下情况不应创建节点： +- 完全无输出的工具调用 +- 纯系统错误（与目标无关，如本地环境问题） +- 重复的相同失败（多次相同错误只保留第一次） + +### 节点合并规则 +以下情况应合并节点： +- 同一工具的多次相似调用（如多次nmap扫描不同端口范围，合并为一个"端口扫描"节点） +- 同一目标的多个相似探测（如多个目录扫描工具，合并为一个"目录扫描"节点） + +### 节点数量控制 +- **完整性优先**：必须包含所有有意义的工具执行和关键步骤，不要为了控制数量而删除重要节点 +- **建议范围**：单目标通常8-15个节点，但如果实际执行步骤较多，可以适当增加（最多20个节点） +- **优先保留**：关键成功步骤、提供线索的失败、发现的漏洞、重要的信息收集步骤 +- **可以合并**：同一工具的多次相似调用（如多次nmap扫描不同端口范围，合并为一个"端口扫描"节点） +- **可以删除**：完全无输出的工具调用、纯系统错误、重复的相同失败（多次相同错误只保留第一次） +- **重要原则**：宁可节点稍多，也不要遗漏关键步骤。攻击链必须能够完整展现渗透测试的完整过程 + +## 边的类型和权重 + +### 边的类型 +- **leads_to**：表示"导致"或"引导到"，用于action→action、target→action + * 例如：端口扫描 → 目录扫描（因为发现了80端口，所以进行目录扫描） +- **discovers**：表示"发现"，**专门用于action→vulnerability** + * 例如：SQL注入测试 → SQL注入漏洞 + * **重要**：所有action→vulnerability的边都必须使用discovers类型，即使多个action都指向同一个vulnerability，也应该统一使用discovers +- **enables**：表示"使能"或"促成"，**仅用于vulnerability→vulnerability、action→action（当后续行动依赖前面结果时）** + * 例如：信息泄露漏洞 → 权限提升漏洞（通过信息泄露获得的信息促成了权限提升） + * **重要**：enables不能用于action→vulnerability，action→vulnerability必须使用discovers + +### 边的权重 +- **权重1-2**：弱关联（如初步探测到进一步探测） +- **权重3-4**：中等关联（如发现端口到服务识别） +- **权重5-7**：强关联（如发现漏洞、关键信息泄露） +- **权重8-10**：极强关联（如漏洞利用成功、权限提升） + +### DAG结构要求（有向无环图） +**关键：必须确保生成的是真正的DAG（有向无环图），不能有任何循环。** + +- **节点编号规则**：节点id从"node_1"开始递增（node_1, node_2, node_3...） +- **边的方向规则**：所有边的source节点id必须严格小于target节点id（source < target），这是确保无环的关键 + * 例如：node_1 → node_2 ✓（正确） + * 例如：node_2 → node_1 ✗（错误，会形成环） + * 例如：node_3 → node_5 ✓（正确） +- **无环验证**：在输出JSON前，必须检查所有边，确保没有任何一条边的source >= target +- **无孤立节点**：确保每个节点至少有一条边连接（除了可能的根节点） +- **DAG结构特点**： + * 一个节点可以有多个后续节点（分支），例如：node_2（端口扫描）可以同时连接到node_3、node_4、node_5等多个节点 + * 多个节点可以汇聚到一个节点（汇聚），例如：node_3、node_4、node_5都指向node_6（漏洞节点） + * 避免将所有节点连成一条线，应该根据实际的并行测试和分支探索构建DAG结构 +- **拓扑排序验证**：如果按照节点id从小到大排序，所有边都应该从左指向右（从上指向下），这样就能保证无环 + +## 攻击链逻辑连贯性要求 + +构建的攻击链应该能够回答以下问题： +1. **起点**：测试从哪里开始？（target节点） +2. **探索过程**：如何逐步收集信息？（action节点序列） +3. **失败与调整**：遇到障碍时如何调整策略？（failed_insight节点） +4. **关键发现**：发现了哪些重要信息？（action的findings） +5. **漏洞确认**：如何确认漏洞存在？（action→vulnerability） +6. **攻击路径**：完整的攻击路径是什么？（从target到vulnerability的路径） + +## 当前任务 ReAct 轨迹（含工具执行；助手结论见轨迹末尾 assistant） + +%s +%s + +## 输出格式 + +严格按照以下JSON格式输出，不要添加任何其他文字： + +**重要：示例展示的是树状结构，注意node_2（端口扫描）同时连接到多个后续节点（node_3、node_4），形成分支结构。** + +{ + "nodes": [ + { + "id": "node_1", + "type": "target", + "label": "测试目标: example.com", + "risk_score": 40, + "metadata": { + "target": "example.com" + } + }, + { + "id": "node_2", + "type": "action", + "label": "扫描端口发现80/443/8080", + "risk_score": 0, + "metadata": { + "tool_name": "nmap", + "tool_intent": "端口扫描", + "ai_analysis": "使用nmap对目标进行端口扫描，发现80、443、8080端口开放。80端口运行HTTP服务，443端口运行HTTPS服务，8080端口可能为管理后台。这些开放端口为后续Web应用测试提供了入口。", + "findings": ["80端口开放", "443端口开放", "8080端口开放", "HTTP服务为Apache 2.4"] + } + }, + { + "id": "node_3", + "type": "action", + "label": "目录扫描发现/admin后台", + "risk_score": 0, + "metadata": { + "tool_name": "dirsearch", + "tool_intent": "目录扫描", + "ai_analysis": "使用dirsearch对目标进行目录扫描，发现/admin目录存在且可访问。该目录可能为管理后台，是重要的测试目标。", + "findings": ["/admin目录存在", "返回200状态码", "疑似管理后台"] + } + }, + { + "id": "node_4", + "type": "action", + "label": "识别Web服务为Apache 2.4", + "risk_score": 0, + "metadata": { + "tool_name": "whatweb", + "tool_intent": "Web服务识别", + "ai_analysis": "识别出目标运行Apache 2.4服务器，这为后续的漏洞测试提供了重要信息。", + "findings": ["Apache 2.4", "PHP版本信息"] + } + }, + { + "id": "node_5", + "type": "action", + "label": "尝试SQL注入（被WAF拦截）", + "risk_score": 0, + "metadata": { + "tool_name": "sqlmap", + "tool_intent": "SQL注入检测", + "ai_analysis": "对/login.php进行SQL注入测试时被WAF拦截，返回403错误。错误信息显示检测到Cloudflare防护。这表明目标部署了WAF，需要调整测试策略。", + "findings": ["WAF拦截", "返回403", "检测到Cloudflare", "目标部署WAF"], + "status": "failed_insight" + } + }, + { + "id": "node_6", + "type": "vulnerability", + "label": "SQL注入漏洞", + "risk_score": 85, + "metadata": { + "vulnerability_type": "SQL注入", + "description": "在/admin/login.php的username参数发现SQL注入漏洞，可通过注入payload绕过登录验证，直接获取管理员权限。漏洞返回数据库错误信息，确认存在注入点。", + "severity": "high", + "location": "/admin/login.php?username=" + } + } + ], + "edges": [ + { + "source": "node_1", + "target": "node_2", + "type": "leads_to", + "weight": 3 + }, + { + "source": "node_2", + "target": "node_3", + "type": "leads_to", + "weight": 4 + }, + { + "source": "node_2", + "target": "node_4", + "type": "leads_to", + "weight": 3 + }, + { + "source": "node_3", + "target": "node_5", + "type": "leads_to", + "weight": 4 + }, + { + "source": "node_5", + "target": "node_6", + "type": "discovers", + "weight": 7 + } + ] +} + +## 重要提醒 + +1. **严禁杜撰**：只使用ReAct输入中实际执行的工具和实际返回的结果。如无实际数据，返回空的nodes和edges数组。 +2. **DAG结构必须**：必须构建真正的DAG（有向无环图），不能有任何循环。所有边的source节点id必须严格小于target节点id（source < target）。 +3. **拓扑顺序**：节点应该按照逻辑顺序编号，target节点通常是node_1，后续的action节点按执行顺序递增，vulnerability节点在最后。 +4. **完整性优先**：必须包含所有有意义的工具执行和关键步骤，不要为了控制节点数量而删除重要节点。攻击链必须能够完整展现从目标识别到漏洞发现的完整过程。 +5. **逻辑连贯**：确保攻击链能够讲述一个完整、连贯的渗透测试故事，包括所有关键步骤和决策点。 +6. **教育价值**：优先保留有教育意义的节点，帮助学习者理解渗透测试思维和完整流程。 +7. **准确性**：所有节点信息必须基于实际数据，不要推测或假设。 +8. **完整性检查**：确保每个节点都有必要的metadata字段，每条边都有正确的source和target，没有孤立节点，没有循环。 +9. **不要过度精简**：如果实际执行步骤较多，可以适当增加节点数量（最多20个），确保不遗漏关键步骤。 +10. **输出前验证**：在输出JSON前，必须验证所有边都满足source < target的条件，确保DAG结构正确。 + +现在开始分析并构建攻击链：`, reactInput, assistantOutSection(modelOutput)) +} + +func assistantOutSection(modelOutput string) string { + modelOutput = strings.TrimSpace(modelOutput) + if modelOutput == "" { + return "" + } + return "\n## 助手结论（补充）\n\n" + modelOutput + "\n" +} + +// saveChain 保存攻击链到数据库 +func (b *Builder) saveChain(conversationID string, nodes []Node, edges []Edge) error { + // 先删除旧的攻击链数据 + if err := b.db.DeleteAttackChain(conversationID); err != nil { + b.logger.Warn("删除旧攻击链失败", zap.Error(err)) + } + + for _, node := range nodes { + metadataJSON, _ := json.Marshal(node.Metadata) + if err := b.db.SaveAttackChainNode(conversationID, node.ID, node.Type, node.Label, "", string(metadataJSON), node.RiskScore); err != nil { + b.logger.Warn("保存攻击链节点失败", zap.String("nodeId", node.ID), zap.Error(err)) + } + } + + // 保存边 + for _, edge := range edges { + if err := b.db.SaveAttackChainEdge(conversationID, edge.ID, edge.Source, edge.Target, edge.Type, edge.Weight); err != nil { + b.logger.Warn("保存攻击链边失败", zap.String("edgeId", edge.ID), zap.Error(err)) + } + } + + return nil +} + +// LoadChainFromDatabase 从数据库加载攻击链 +func (b *Builder) LoadChainFromDatabase(conversationID string) (*Chain, error) { + nodes, err := b.db.LoadAttackChainNodes(conversationID) + if err != nil { + return nil, fmt.Errorf("加载攻击链节点失败: %w", err) + } + + edges, err := b.db.LoadAttackChainEdges(conversationID) + if err != nil { + return nil, fmt.Errorf("加载攻击链边失败: %w", err) + } + + return &Chain{ + Nodes: nodes, + Edges: edges, + }, nil +} + +// callAIForChainGeneration 调用AI生成攻击链 +func (b *Builder) callAIForChainGeneration(ctx context.Context, prompt string) (string, error) { + requestBody := map[string]interface{}{ + "model": b.openAIConfig.Model, + "messages": []map[string]interface{}{ + { + "role": "system", + "content": "你是一个专业的安全测试分析师，擅长构建攻击链图。请严格按照JSON格式返回攻击链数据。", + }, + { + "role": "user", + "content": prompt, + }, + }, + "temperature": 0.3, + "max_completion_tokens": attackChainMaxCompletionTokens(b.maxTokens), + } + + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + if b.openAIClient == nil { + return "", fmt.Errorf("OpenAI客户端未初始化") + } + if err := b.openAIClient.ChatCompletion(ctx, requestBody, &apiResponse); err != nil { + var apiErr *openai.APIError + if errors.As(err, &apiErr) { + bodyStr := strings.ToLower(apiErr.Body) + if strings.Contains(bodyStr, "context") || strings.Contains(bodyStr, "length") || strings.Contains(bodyStr, "too long") { + return "", fmt.Errorf("context length exceeded") + } + } else if strings.Contains(strings.ToLower(err.Error()), "context") || strings.Contains(strings.ToLower(err.Error()), "length") { + return "", fmt.Errorf("context length exceeded") + } + return "", fmt.Errorf("请求失败: %w", err) + } + + if len(apiResponse.Choices) == 0 { + return "", fmt.Errorf("API未返回有效响应") + } + + content := strings.TrimSpace(apiResponse.Choices[0].Message.Content) + // 尝试提取JSON（可能包含markdown代码块） + content = strings.TrimPrefix(content, "```json") + content = strings.TrimPrefix(content, "```") + content = strings.TrimSuffix(content, "```") + content = strings.TrimSpace(content) + + return content, nil +} + +// ChainJSON 攻击链JSON结构 +type ChainJSON struct { + Nodes []struct { + ID string `json:"id"` + Type string `json:"type"` + Label string `json:"label"` + RiskScore int `json:"risk_score"` + Metadata map[string]interface{} `json:"metadata"` + } `json:"nodes"` + Edges []struct { + Source string `json:"source"` + Target string `json:"target"` + Type string `json:"type"` + Weight int `json:"weight"` + } `json:"edges"` +} + +// parseChainJSON 解析攻击链JSON +func (b *Builder) parseChainJSON(chainJSON string) (*Chain, error) { + var chainData ChainJSON + if err := json.Unmarshal([]byte(chainJSON), &chainData); err != nil { + return nil, fmt.Errorf("解析JSON失败: %w", err) + } + + // 创建节点ID映射（AI返回的ID -> 新的UUID） + nodeIDMap := make(map[string]string) + + // 转换为Chain结构 + nodes := make([]Node, 0, len(chainData.Nodes)) + for _, n := range chainData.Nodes { + // 生成新的UUID节点ID + newNodeID := fmt.Sprintf("node_%s", uuid.New().String()) + nodeIDMap[n.ID] = newNodeID + + node := Node{ + ID: newNodeID, + Type: n.Type, + Label: n.Label, + RiskScore: n.RiskScore, + Metadata: n.Metadata, + } + if node.Metadata == nil { + node.Metadata = make(map[string]interface{}) + } + nodes = append(nodes, node) + } + + // 转换边 + edges := make([]Edge, 0, len(chainData.Edges)) + for _, e := range chainData.Edges { + sourceID, ok := nodeIDMap[e.Source] + if !ok { + continue + } + targetID, ok := nodeIDMap[e.Target] + if !ok { + continue + } + + // 生成边的ID（前端需要） + edgeID := fmt.Sprintf("edge_%s", uuid.New().String()) + + edges = append(edges, Edge{ + ID: edgeID, + Source: sourceID, + Target: targetID, + Type: e.Type, + Weight: e.Weight, + }) + } + + return &Chain{ + Nodes: nodes, + Edges: edges, + }, nil +} + +// 以下所有方法已不再使用，已删除以简化代码 diff --git a/internal/attackchain/promote_project.go b/internal/attackchain/promote_project.go new file mode 100644 index 00000000..d8a9cd80 --- /dev/null +++ b/internal/attackchain/promote_project.go @@ -0,0 +1,203 @@ +package attackchain + +import ( + "fmt" + "regexp" + "strings" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/project" + + "github.com/google/uuid" +) + +var promoteSlugSanitizer = regexp.MustCompile(`[^a-z0-9._/-]+`) + +// PromoteToProjectResult 攻击链沉淀结果。 +type PromoteToProjectResult struct { + FactsCreated int `json:"facts_created"` + FactsUpdated int `json:"facts_updated"` + EdgesCreated int `json:"edges_created"` + FactKeys []string `json:"fact_keys"` + Graph *database.ProjectFactGraph `json:"graph,omitempty"` +} + +// PromoteToProject 将对话攻击链沉淀为项目事实与边。 +func PromoteToProject(db *database.DB, projectID, conversationID string) (*PromoteToProjectResult, error) { + if db == nil { + return nil, fmt.Errorf("database 未初始化") + } + projectID = strings.TrimSpace(projectID) + conversationID = strings.TrimSpace(conversationID) + if projectID == "" || conversationID == "" { + return nil, fmt.Errorf("project_id 与 conversation_id 必填") + } + if _, err := db.GetProject(projectID); err != nil { + return nil, fmt.Errorf("项目不存在") + } + conv, err := db.GetConversation(conversationID) + if err != nil { + return nil, fmt.Errorf("对话不存在") + } + if pid := strings.TrimSpace(conv.ProjectID); pid != "" && pid != projectID { + return nil, fmt.Errorf("对话已绑定其他项目") + } + + nodes, err := db.LoadAttackChainNodes(conversationID) + if err != nil { + return nil, err + } + edges, err := db.LoadAttackChainEdges(conversationID) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, fmt.Errorf("该对话尚无攻击链，请先在对话中生成攻击链") + } + + res := &PromoteToProjectResult{} + nodeToKey := make(map[string]string, len(nodes)) + usedKeys := map[string]int{} + + for _, node := range nodes { + key := allocatePromoteFactKey(node, usedKeys) + nodeToKey[node.ID] = key + category := mapPromoteNodeCategory(node.Type) + existing, getErr := db.GetProjectFactByKey(projectID, key) + f := &database.ProjectFact{ + ProjectID: projectID, + FactKey: key, + Category: category, + Summary: strings.TrimSpace(node.Label), + Body: formatPromotedFactBody(node, conversationID), + Confidence: "tentative", + SourceConversationID: conversationID, + } + if getErr == nil && existing != nil { + f.ID = existing.ID + f.CreatedAt = existing.CreatedAt + if strings.TrimSpace(f.Summary) == "" { + f.Summary = existing.Summary + } + if _, err := db.UpsertProjectFact(f); err != nil { + return nil, err + } + res.FactsUpdated++ + } else { + if _, err := db.UpsertProjectFact(f); err != nil { + return nil, err + } + res.FactsCreated++ + } + res.FactKeys = append(res.FactKeys, key) + } + + for _, edge := range edges { + srcKey, ok1 := nodeToKey[edge.Source] + tgtKey, ok2 := nodeToKey[edge.Target] + if !ok1 || !ok2 || srcKey == tgtKey { + continue + } + edgeType := mapPromoteEdgeType(edge.Type) + incoming, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey) + merged := project.MergeLinkFromInputsUnique(promoteFromEdgeInputsFromDB(incoming), []database.ProjectFactEdgeFromInput{{From: srcKey, Type: edgeType}}) + if err := db.ReplaceIncomingProjectFactEdges(projectID, tgtKey, merged); err != nil { + return nil, err + } + res.EdgesCreated++ + if fact, err := db.GetProjectFactByKey(projectID, tgtKey); err == nil { + in, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey) + fact.Body = project.SyncBodyLinksSection(fact.Body, in) + _, _ = db.UpsertProjectFact(fact) + } + } + + graph, _ := project.BuildProjectFactGraph(db, projectID, "full", true) + res.Graph = graph + return res, nil +} + +func promoteFromEdgeInputsFromDB(edges []*database.ProjectFactEdge) []database.ProjectFactEdgeFromInput { + out := make([]database.ProjectFactEdgeFromInput, 0, len(edges)) + for _, e := range edges { + out = append(out, database.ProjectFactEdgeFromInput{From: e.SourceFactKey, Type: e.EdgeType, Confidence: e.Confidence}) + } + return out +} + +func mapPromoteNodeCategory(nodeType string) string { + switch strings.ToLower(strings.TrimSpace(nodeType)) { + case "target": + return project.FactCategoryTarget + case "vulnerability": + return project.FactCategoryFinding + case "action": + return project.FactCategoryChain + default: + return project.FactCategoryNote + } +} + +func mapPromoteEdgeType(t string) string { + switch strings.ToLower(strings.TrimSpace(t)) { + case "discovers", "discovered_on", "targets": + return "discovered_on" + case "exploits": + return "exploits" + case "enables": + return "enables" + case "depends_on": + return "depends_on" + default: + return "leads_to" + } +} + +func allocatePromoteFactKey(node Node, used map[string]int) string { + prefix := "chain/" + switch strings.ToLower(strings.TrimSpace(node.Type)) { + case "target": + prefix = "target/" + case "vulnerability": + prefix = "finding/" + case "action": + prefix = "chain/" + } + base := promoteSlugify(node.Label) + if base == "" { + base = promoteSlugify(node.ID) + } + if base == "" { + base = uuid.New().String()[:8] + } + key := prefix + base + if n, ok := used[key]; ok { + n++ + used[key] = n + key = fmt.Sprintf("%s-%d", key, n) + } else { + used[key] = 1 + } + return key +} + +func promoteSlugify(s string) string { + s = strings.ToLower(strings.TrimSpace(s)) + s = strings.NewReplacer(" ", "-", "—", "-", "–", "-", "/", "-").Replace(s) + s = promoteSlugSanitizer.ReplaceAllString(s, "-") + s = strings.Trim(s, "-") + if len(s) > 64 { + s = s[:64] + } + return s +} + +func formatPromotedFactBody(node Node, conversationID string) string { + var b strings.Builder + b.WriteString("## 来源\n") + b.WriteString(fmt.Sprintf("- 对话攻击链沉淀\n- source_conversation_id: %s\n- node_id: %s\n- node_type: %s\n\n", conversationID, node.ID, node.Type)) + b.WriteString("## 摘要\n") + b.WriteString(strings.TrimSpace(node.Label)) + b.WriteString("\n\n## 关联\n- 结构化关系边（自动同步）:\n （见项目攻击路径图）\n") + return b.String() +} diff --git a/internal/attackchain/truncate.go b/internal/attackchain/truncate.go new file mode 100644 index 00000000..ba379b3b --- /dev/null +++ b/internal/attackchain/truncate.go @@ -0,0 +1,248 @@ +package attackchain + +import ( + "strings" + "unicode/utf8" + + "go.uber.org/zap" +) + +const ( + attackChainTruncationMarker = "\n\n...[攻击链输入已截断 / attack chain input truncated]...\n\n" + attackChainSystemReserve = 256 + attackChainSafetyReserve = 2048 +) + +// attackChainMaxCompletionTokens 为攻击链 JSON 输出预留的 completion token 上限。 +func attackChainMaxCompletionTokens(maxTotal int) int { + const capTokens = 16384 + if maxTotal <= 0 { + return 8192 + } + v := maxTotal / 8 + if v < 4096 { + v = 4096 + } + if v > capTokens { + v = capTokens + } + return v +} + +func (b *Builder) modelName() string { + if b.openAIConfig != nil && b.openAIConfig.Model != "" { + return b.openAIConfig.Model + } + return "gpt-4" +} + +func (b *Builder) countTokens(text string) int { + if text == "" { + return 0 + } + n, err := b.tokenCounter.Count(b.modelName(), text) + if err != nil { + return utf8.RuneCountInString(text) / 4 + } + return n +} + +// attackChainPayloadTokenBudget 计算 reactInput + modelOutput 可用的 token 预算。 +func (b *Builder) attackChainPayloadTokenBudget() int { + maxTotal := b.maxTokens + if maxTotal <= 0 { + maxTotal = 100000 + } + templateTok := b.countTokens(b.buildSimplePrompt("", "")) + completion := attackChainMaxCompletionTokens(maxTotal) + reserve := templateTok + attackChainSystemReserve + completion + attackChainSafetyReserve + budget := maxTotal - reserve + minBudget := maxTotal * 35 / 100 + if budget < minBudget { + budget = minBudget + } + if budget < 4096 { + budget = 4096 + } + return budget +} + +// fitAttackChainPayload 在构建最终 prompt 前压缩 ReAct 轨迹与模型输出，避免超出模型上下文。 +func (b *Builder) fitAttackChainPayload(reactInput, modelOutput string) (string, string, bool) { + budget := b.attackChainPayloadTokenBudget() + modelBudget := budget * 15 / 100 + if modelBudget < 512 { + modelBudget = 512 + } + reactBudget := budget - modelBudget + + origReactTok := b.countTokens(reactInput) + origModelTok := b.countTokens(modelOutput) + truncated := false + + outModel := modelOutput + if origModelTok > modelBudget { + outModel = truncateTextByTokens(b, modelOutput, modelBudget) + truncated = true + } + + outReact := reactInput + perToolLimits := []int{12000, 6000, 3000, 1500, 800} + for _, lim := range perToolLimits { + compact := compactFormattedToolBodies(outReact, lim) + if compact != outReact { + outReact = compact + truncated = true + } + if b.countTokens(outReact) <= reactBudget { + break + } + } + + if b.countTokens(outReact) > reactBudget { + outReact = truncateTextByTokens(b, outReact, reactBudget) + truncated = true + } + + if truncated { + b.logger.Info("攻击链输入已按 token 预算截断", + zap.Int("maxTotalTokens", b.maxTokens), + zap.Int("payloadBudget", budget), + zap.Int("reactBudget", reactBudget), + zap.Int("modelBudget", modelBudget), + zap.Int("reactInputTokensBefore", origReactTok), + zap.Int("reactInputTokensAfter", b.countTokens(outReact)), + zap.Int("modelOutputTokensBefore", origModelTok), + zap.Int("modelOutputTokensAfter", b.countTokens(outModel)), + zap.Int("maxCompletionTokens", attackChainMaxCompletionTokens(b.maxTokens)), + ) + } + + return outReact, outModel, truncated +} + +// compactFormattedToolBodies 缩短格式化 trace 中 [tool] 消息的正文，保留工具头与调用 ID。 +func compactFormattedToolBodies(s string, maxRunesPerBody int) string { + if maxRunesPerBody <= 0 || s == "" { + return s + } + const marker = "[tool]" + var out strings.Builder + remaining := s + changed := false + for { + idx := strings.Index(remaining, marker) + if idx < 0 { + out.WriteString(remaining) + break + } + out.WriteString(remaining[:idx]) + remaining = remaining[idx:] + nl := strings.IndexByte(remaining, '\n') + if nl < 0 { + out.WriteString(remaining) + break + } + header := remaining[:nl+1] + remaining = remaining[nl+1:] + bodyEnd := strings.Index(remaining, "\n\n[") + var body, rest string + if bodyEnd < 0 { + body = remaining + rest = "" + } else { + body = remaining[:bodyEnd] + rest = remaining[bodyEnd:] + } + if runeLen(body) > maxRunesPerBody { + body = truncateRunesWithNotice(body, maxRunesPerBody) + changed = true + } + out.WriteString(header) + out.WriteString(body) + remaining = rest + if rest == "" { + break + } + } + if !changed { + return s + } + return out.String() +} + +func truncateTextByTokens(b *Builder, text string, maxTokens int) string { + if maxTokens <= 0 || text == "" { + return "" + } + if b.countTokens(text) <= maxTokens { + return text + } + markerTok := b.countTokens(attackChainTruncationMarker) + usable := maxTokens - markerTok + if usable < 256 { + usable = maxTokens / 2 + } + headBudget := usable * 60 / 100 + tailBudget := usable - headBudget + head := takeTokensFromStart(b, text, headBudget) + tail := takeTokensFromEnd(b, text, tailBudget) + return head + attackChainTruncationMarker + tail +} + +func takeTokensFromStart(b *Builder, text string, maxTokens int) string { + rs := []rune(text) + if len(rs) == 0 || maxTokens <= 0 { + return "" + } + lo, hi := 0, len(rs) + for lo < hi { + mid := (lo + hi + 1) / 2 + if b.countTokens(string(rs[:mid])) <= maxTokens { + lo = mid + } else { + hi = mid - 1 + } + } + return string(rs[:lo]) +} + +func takeTokensFromEnd(b *Builder, text string, maxTokens int) string { + rs := []rune(text) + if len(rs) == 0 || maxTokens <= 0 { + return "" + } + lo, hi := 0, len(rs) + for lo < hi { + mid := (lo + hi) / 2 + if b.countTokens(string(rs[mid:])) <= maxTokens { + hi = mid + } else { + lo = mid + 1 + } + } + return string(rs[lo:]) +} + +func truncateRunesWithNotice(s string, maxRunes int) string { + rs := []rune(s) + if len(rs) <= maxRunes { + return s + } + const notice = "\n...[工具输出已截断 / tool output truncated]...\n" + noticeRunes := []rune(notice) + keep := maxRunes - len(noticeRunes) + if keep < 200 { + keep = maxRunes * 2 / 3 + } + if keep < 1 { + return notice + } + head := keep * 70 / 100 + tail := keep - head + return string(rs[:head]) + notice + string(rs[len(rs)-tail:]) +} + +func runeLen(s string) int { + return len([]rune(s)) +} diff --git a/internal/attackchain/truncate_test.go b/internal/attackchain/truncate_test.go new file mode 100644 index 00000000..2cb4563c --- /dev/null +++ b/internal/attackchain/truncate_test.go @@ -0,0 +1,63 @@ +package attackchain + +import ( + "strings" + "testing" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func testBuilder(maxTotal int) *Builder { + return &Builder{ + logger: zap.NewNop(), + openAIConfig: &config.OpenAIConfig{Model: "gpt-4"}, + tokenCounter: agent.NewTikTokenCounter(), + maxTokens: maxTotal, + } +} + +func TestCompactFormattedToolBodies(t *testing.T) { + long := strings.Repeat("x", 20000) + in := "[user]: hi\n\n[tool] (tool_call_id: abc):\n" + long + "\n\n[assistant]: done\n" + out := compactFormattedToolBodies(in, 500) + if strings.Contains(out, strings.Repeat("x", 10000)) { + t.Fatal("expected tool body to be truncated") + } + if !strings.Contains(out, "[user]: hi") { + t.Fatal("expected user header preserved") + } + if !strings.Contains(out, "[assistant]: done") { + t.Fatal("expected assistant header preserved") + } +} + +func TestFitAttackChainPayloadWithinBudget(t *testing.T) { + b := testBuilder(32000) + react := strings.Repeat("scan ", 50000) + model := strings.Repeat("result ", 10000) + r, m, truncated := b.fitAttackChainPayload(react, model) + if !truncated { + t.Fatal("expected truncation for large payload") + } + prompt := b.buildSimplePrompt(r, m) + total := b.countTokens(prompt) + attackChainMaxCompletionTokens(b.maxTokens) + attackChainSystemReserve + if total > b.maxTokens+attackChainSafetyReserve { + t.Fatalf("prompt still too large: estimated %d > max %d", total, b.maxTokens) + } + _ = m +} + +func TestAttackChainMaxCompletionTokens(t *testing.T) { + if got := attackChainMaxCompletionTokens(120000); got != 15000 && got != 16384 { + // 120000/8 = 15000 + if got < 4096 || got > 16384 { + t.Fatalf("unexpected completion cap: %d", got) + } + } + if got := attackChainMaxCompletionTokens(0); got != 8192 { + t.Fatalf("expected default 8192, got %d", got) + } +} diff --git a/internal/audit/conversation_create.go b/internal/audit/conversation_create.go new file mode 100644 index 00000000..82e19b54 --- /dev/null +++ b/internal/audit/conversation_create.go @@ -0,0 +1,55 @@ +package audit + +import ( + "strings" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" +) + +// RegisterConversationCreateHook records platform audit rows for every new conversation. +func RegisterConversationCreateHook(s *Service) { + if s == nil { + return + } + database.SetConversationCreateHook(func(conv *database.Conversation, meta database.ConversationCreateMeta) { + detail := map[string]interface{}{ + "title": conv.Title, + "source": meta.Source, + } + if meta.WebShellConnectionID != "" { + detail["webshell_connection_id"] = meta.WebShellConnectionID + } + s.Record(nil, Entry{ + Category: "conversation", + Action: "create", + Result: "success", + Message: "创建对话", + ResourceType: "conversation", + ResourceID: conv.ID, + Detail: detail, + ClientIP: meta.ClientIP, + SessionHint: meta.SessionHint, + }) + }) +} + +// ConversationCreateMeta builds audit metadata for conversation creation. +func ConversationCreateMeta(source string) database.ConversationCreateMeta { + return database.ConversationCreateMeta{Source: strings.TrimSpace(source)} +} + +// ConversationCreateMetaFromGin includes client IP and session hint when available. +func ConversationCreateMetaFromGin(c *gin.Context, source string) database.ConversationCreateMeta { + m := ConversationCreateMeta(source) + if c == nil { + return m + } + m.ClientIP = c.ClientIP() + if token := c.GetString(security.ContextAuthTokenKey); token != "" { + m.SessionHint = sessionHint(token) + } + return m +} diff --git a/internal/audit/meta.go b/internal/audit/meta.go new file mode 100644 index 00000000..33649e0c --- /dev/null +++ b/internal/audit/meta.go @@ -0,0 +1,9 @@ +package audit + +// RetentionDays returns configured retention; 0 means keep forever. +func (s *Service) RetentionDays() int { + if s == nil || s.cfg == nil { + return 0 + } + return s.cfg.Audit.RetentionDaysEffective() +} diff --git a/internal/audit/record.go b/internal/audit/record.go new file mode 100644 index 00000000..b1c1ad40 --- /dev/null +++ b/internal/audit/record.go @@ -0,0 +1,29 @@ +package audit + +import "github.com/gin-gonic/gin" + +// RecordAction writes a platform audit row with common defaults. +func (s *Service) RecordAction(c *gin.Context, category, action, result, message, resourceType, resourceID string, detail map[string]interface{}) { + if s == nil { + return + } + s.Record(c, Entry{ + Category: category, + Action: action, + Result: result, + Message: message, + ResourceType: resourceType, + ResourceID: resourceID, + Detail: detail, + }) +} + +// RecordOK is a shorthand for successful operations. +func (s *Service) RecordOK(c *gin.Context, category, action, message, resourceType, resourceID string, detail map[string]interface{}) { + s.RecordAction(c, category, action, "success", message, resourceType, resourceID, detail) +} + +// RecordFail is a shorthand for failed operations. +func (s *Service) RecordFail(c *gin.Context, category, action, message string, detail map[string]interface{}) { + s.RecordAction(c, category, action, "failure", message, "", "", detail) +} diff --git a/internal/audit/resource_availability.go b/internal/audit/resource_availability.go new file mode 100644 index 00000000..3b22871f --- /dev/null +++ b/internal/audit/resource_availability.go @@ -0,0 +1,86 @@ +package audit + +import ( + "strings" + + "cyberstrike-ai/internal/database" +) + +var auditActionsResourceRemoved = map[string]bool{ + "delete": true, + "item_delete": true, + "connection_delete": true, + "listener_delete": true, + "session_delete": true, + "task_delete": true, + "execution_delete": true, + "execution_delete_batch": true, + "delete_queue": true, + "delete_batch_task": true, + "markdown_delete": true, +} + +// ApplyResourceAvailability sets log.ResourceAvailable when the linked resource can be checked. +func ApplyResourceAvailability(db *database.DB, log *database.AuditLog) { + if log == nil || strings.TrimSpace(log.ResourceID) == "" { + return + } + if auditActionsResourceRemoved[log.Action] { + f := false + log.ResourceAvailable = &f + return + } + if db == nil { + return + } + available, known := resourceStillExists(db, log.ResourceType, log.ResourceID) + if known { + log.ResourceAvailable = &available + } +} + +func resourceStillExists(db *database.DB, resourceType, resourceID string) (bool, bool) { + resourceID = strings.TrimSpace(resourceID) + if resourceID == "" { + return false, false + } + t := strings.TrimSpace(resourceType) + if t == "" { + if len(resourceID) > 8 && !strings.HasPrefix(resourceID, "c2_") { + t = "conversation" + } else { + return false, false + } + } + switch t { + case "conversation": + ok, err := db.ConversationExists(resourceID) + return ok, err == nil + case "vulnerability": + _, err := db.GetVulnerability(resourceID) + if err != nil { + return false, strings.Contains(err.Error(), "不存在") + } + return true, true + case "batch_queue": + _, err := db.GetBatchQueue(resourceID) + return err == nil, true + case "c2_listener": + _, err := db.GetC2Listener(resourceID) + return err == nil, true + case "c2_session": + _, err := db.GetC2Session(resourceID) + return err == nil, true + case "c2_task": + _, err := db.GetC2Task(resourceID) + return err == nil, true + case "webshell_connection": + c, err := db.GetWebshellConnection(resourceID) + return err == nil && c != nil, true + case "tool_execution": + _, err := db.GetToolExecution(resourceID) + return err == nil, true + default: + return false, false + } +} diff --git a/internal/audit/retention.go b/internal/audit/retention.go new file mode 100644 index 00000000..f882595c --- /dev/null +++ b/internal/audit/retention.go @@ -0,0 +1,27 @@ +package audit + +import ( + "time" + + "go.uber.org/zap" +) + +// auditRetentionPurgeInterval is how often PurgeExpired runs while the process is up (startup also purges once). +const auditRetentionPurgeInterval = time.Hour + +// StartRetentionLoop periodically purges expired audit rows. +func StartRetentionLoop(s *Service, logger *zap.Logger) { + if s == nil { + return + } + go func() { + ticker := time.NewTicker(auditRetentionPurgeInterval) + defer ticker.Stop() + for range ticker.C { + s.PurgeExpired() + if logger != nil { + logger.Debug("audit retention tick completed") + } + } + }() +} diff --git a/internal/audit/sanitize.go b/internal/audit/sanitize.go new file mode 100644 index 00000000..34f2b439 --- /dev/null +++ b/internal/audit/sanitize.go @@ -0,0 +1,58 @@ +package audit + +import ( + "encoding/json" + "strings" +) + +var sensitiveKeySubstrings = []string{ + "password", "api_key", "apikey", "secret", "token", "authorization", + "credential", "private_key", "access_key", +} + +// SanitizeDetail redacts sensitive keys and truncates serialized size. +func SanitizeDetail(detail map[string]interface{}, maxBytes int) map[string]interface{} { + if detail == nil { + return nil + } + if maxBytes <= 0 { + maxBytes = 8192 + } + out := sanitizeValue("", detail) + if m, ok := out.(map[string]interface{}); ok { + b, _ := json.Marshal(m) + if len(b) > maxBytes { + return map[string]interface{}{ + "_truncated": true, + "_preview": string(b[:maxBytes]), + } + } + return m + } + return map[string]interface{}{"value": out} +} + +func sanitizeValue(key string, v interface{}) interface{} { + kl := strings.ToLower(key) + for _, sub := range sensitiveKeySubstrings { + if strings.Contains(kl, sub) { + return "***" + } + } + switch t := v.(type) { + case map[string]interface{}: + m := make(map[string]interface{}, len(t)) + for k, val := range t { + m[k] = sanitizeValue(k, val) + } + return m + case []interface{}: + arr := make([]interface{}, len(t)) + for i, val := range t { + arr[i] = sanitizeValue(key, val) + } + return arr + default: + return v + } +} diff --git a/internal/audit/service.go b/internal/audit/service.go new file mode 100644 index 00000000..60ea11b3 --- /dev/null +++ b/internal/audit/service.go @@ -0,0 +1,177 @@ +package audit + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Service persists platform audit logs. +type Service struct { + db *database.DB + cfg *config.Config + logger *zap.Logger + failThrottle *failureThrottle +} + +// NewService creates an audit service. +func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service { + return &Service{ + db: db, + cfg: cfg, + logger: logger, + failThrottle: newFailureThrottle(), + } +} + +// Enabled reports whether audit persistence is on. +func (s *Service) Enabled() bool { + if s == nil || s.cfg == nil { + return false + } + return s.cfg.Audit.EnabledEffective() +} + +// Record writes one audit row from a Gin request context. +func (s *Service) Record(c *gin.Context, e Entry) { + if s == nil || !s.Enabled() || s.db == nil { + return + } + if strings.TrimSpace(e.Category) == "" || strings.TrimSpace(e.Action) == "" { + return + } + if e.Result == "failure" && !s.allowFailureAudit(c, e) { + return + } + if strings.TrimSpace(e.Result) == "" { + e.Result = "success" + } + if strings.TrimSpace(e.Level) == "" { + if e.Result == "failure" { + e.Level = "warn" + } else { + e.Level = "info" + } + } + if strings.TrimSpace(e.Actor) == "" { + if c != nil { + e.Actor = strings.TrimSpace(c.GetString(security.ContextUsernameKey)) + } + if e.Actor == "" { + e.Actor = "admin" + } + } + maxDetail := s.cfg.Audit.MaxDetailBytesEffective() + detail := SanitizeDetail(e.Detail, maxDetail) + + sessionHintVal := e.SessionHint + if sessionHintVal == "" && c != nil { + if token := c.GetString(security.ContextAuthTokenKey); token != "" { + sessionHintVal = sessionHint(token) + } + } + clientIPVal := e.ClientIP + if clientIPVal == "" { + clientIPVal = clientIP(c) + } + + row := &database.AuditLog{ + ID: "audit_" + strings.ReplaceAll(uuid.New().String(), "-", ""), + CreatedAt: time.Now(), + Level: e.Level, + Category: e.Category, + Action: e.Action, + Result: e.Result, + Actor: e.Actor, + SessionHint: sessionHintVal, + ClientIP: clientIPVal, + UserAgent: userAgent(c), + ResourceType: e.ResourceType, + ResourceID: e.ResourceID, + Message: e.Message, + Detail: detail, + } + if err := s.db.AppendAuditLog(row); err != nil && s.logger != nil { + s.logger.Warn("写入审计日志失败", + zap.String("action", e.Action), + zap.Error(err), + ) + } +} + +// RecordSystem writes an audit row without HTTP context (e.g. retention cleanup). +func (s *Service) RecordSystem(e Entry) { + s.Record(nil, e) +} + +// PurgeExpired deletes rows older than retention_days when configured. +func (s *Service) PurgeExpired() { + if s == nil || s.db == nil || s.cfg == nil { + return + } + days := s.cfg.Audit.RetentionDaysEffective() + if days <= 0 { + return + } + cutoff := time.Now().AddDate(0, 0, -days) + n, err := s.db.DeleteAuditLogsBefore(cutoff) + if err != nil { + if s.logger != nil { + s.logger.Warn("清理过期审计日志失败", zap.Error(err)) + } + return + } + if n > 0 && s.logger != nil { + s.logger.Info("已清理过期审计日志", zap.Int64("deleted", n)) + } +} + +// HintFromToken returns a short stable hash prefix for a session token. +func HintFromToken(token string) string { + return sessionHint(token) +} + +func sessionHint(token string) string { + token = strings.TrimSpace(token) + if token == "" { + return "" + } + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:4]) +} + +func (s *Service) allowFailureAudit(c *gin.Context, e Entry) bool { + if !isAuthFailureThrottled(e.Category, e.Action) { + return true + } + cooldown := time.Duration(s.cfg.Audit.AuthFailureCooldownEffective()) * time.Second + key := authFailureThrottleKey(e.Category, e.Action, clientIP(c)) + return s.failThrottle.allow(key, cooldown) +} + +func clientIP(c *gin.Context) string { + if c == nil { + return "" + } + return c.ClientIP() +} + +func userAgent(c *gin.Context) string { + if c == nil { + return "" + } + ua := c.GetHeader("User-Agent") + if len(ua) > 512 { + return ua[:512] + } + return ua +} diff --git a/internal/audit/throttle.go b/internal/audit/throttle.go new file mode 100644 index 00000000..7364e07d --- /dev/null +++ b/internal/audit/throttle.go @@ -0,0 +1,55 @@ +package audit + +import ( + "sync" + "time" +) + +// failureThrottle deduplicates high-frequency failure audit rows (e.g. wrong password). +type failureThrottle struct { + mu sync.Mutex + last map[string]time.Time +} + +func newFailureThrottle() *failureThrottle { + return &failureThrottle{last: make(map[string]time.Time)} +} + +// allow reports whether a row with the given key may be written now. +func (t *failureThrottle) allow(key string, cooldown time.Duration) bool { + if t == nil || cooldown <= 0 || key == "" { + return true + } + now := time.Now() + t.mu.Lock() + defer t.mu.Unlock() + if prev, ok := t.last[key]; ok && now.Sub(prev) < cooldown { + return false + } + t.last[key] = now + if len(t.last) > 4096 { + for k, ts := range t.last { + if now.Sub(ts) > cooldown*2 { + delete(t.last, k) + } + } + } + return true +} + +// authFailureThrottleKey builds a per-IP key for auth failure deduplication. +func authFailureThrottleKey(category, action, clientIP string) string { + return category + ":" + action + ":" + clientIP +} + +func isAuthFailureThrottled(category, action string) bool { + if category != "auth" { + return false + } + switch action { + case "login", "change_password": + return true + default: + return false + } +} diff --git a/internal/audit/types.go b/internal/audit/types.go new file mode 100644 index 00000000..ff83ea58 --- /dev/null +++ b/internal/audit/types.go @@ -0,0 +1,16 @@ +package audit + +// Entry describes one platform audit record (not chat/tool execution bodies). +type Entry struct { + Level string + Category string + Action string + Result string // success | failure + Actor string + SessionHint string + ResourceType string + ResourceID string + Message string + Detail map[string]interface{} + ClientIP string // optional when c is nil (robot, batch, DB hook) +} diff --git a/internal/einomcp/holder.go b/internal/einomcp/holder.go new file mode 100644 index 00000000..fe56b442 --- /dev/null +++ b/internal/einomcp/holder.go @@ -0,0 +1,21 @@ +package einomcp + +import "sync" + +// ConversationHolder 在每次 DeepAgent 运行前写入会话 ID，供 MCP 工具桥接使用。 +type ConversationHolder struct { + mu sync.RWMutex + id string +} + +func (h *ConversationHolder) Set(id string) { + h.mu.Lock() + h.id = id + h.mu.Unlock() +} + +func (h *ConversationHolder) Get() string { + h.mu.RLock() + defer h.mu.RUnlock() + return h.id +} diff --git a/internal/einomcp/mcp_tools.go b/internal/einomcp/mcp_tools.go new file mode 100644 index 00000000..edff81b4 --- /dev/null +++ b/internal/einomcp/mcp_tools.go @@ -0,0 +1,214 @@ +package einomcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/security" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + "github.com/eino-contrib/jsonschema" +) + +// ExecutionRecorder 可选，在 MCP 工具成功返回且带有 execution id 时回调（用于汇总 mcpExecutionIds）。 +// toolCallID 来自 Eino compose.GetToolCallID，用于与 reduction 后的展示结果关联。 +type ExecutionRecorder func(executionID, toolCallID string) + +// ToolErrorPrefix 用于把内部 MCP 执行结果中的 IsError 标记传递到多代理上层。 +// Eino 工具通道目前只支持返回字符串，因此通过前缀标识，随后在多代理 runner 中解析为 success/isError。 +const ToolErrorPrefix = "__CYBERSTRIKE_AI_TOOL_ERROR__\n" + +// ToolsFromDefinitions 将单 Agent 使用的 OpenAI 风格工具定义转为 Eino InvokableTool，执行时走 Agent 的 MCP 路径。 +// invokeNotify 可选：与 runEinoADKAgentLoop 共享，在 InvokableRun 返回时触发 UI 与 pending 清理（与 ADK Tool 事件去重）。 +// einoAgentName 为该套工具所属 ChatModelAgent 的 Name（主代理或子代理 id），用于 SSE 上的 einoAgent 字段。 +func ToolsFromDefinitions( + ag *agent.Agent, + holder *ConversationHolder, + defs []agent.Tool, + rec ExecutionRecorder, + toolOutputChunk func(toolName, toolCallID, chunk string), + invokeNotify *ToolInvokeNotifyHolder, + einoAgentName string, +) ([]tool.BaseTool, error) { + out := make([]tool.BaseTool, 0, len(defs)) + for _, d := range defs { + if d.Type != "function" || d.Function.Name == "" { + continue + } + info, err := toolInfoFromDefinition(d) + if err != nil { + return nil, fmt.Errorf("tool %q: %w", d.Function.Name, err) + } + out = append(out, &mcpBridgeTool{ + info: info, + name: d.Function.Name, + agent: ag, + holder: holder, + record: rec, + chunk: toolOutputChunk, + invokeNotify: invokeNotify, + einoAgentName: strings.TrimSpace(einoAgentName), + }) + } + return out, nil +} + +func toolInfoFromDefinition(d agent.Tool) (*schema.ToolInfo, error) { + fn := d.Function + raw, err := json.Marshal(fn.Parameters) + if err != nil { + return nil, err + } + var js jsonschema.Schema + if len(raw) > 0 && string(raw) != "null" && string(raw) != "{}" { + if err := json.Unmarshal(raw, &js); err != nil { + return nil, err + } + } + if js.Type == "" { + js.Type = string(schema.Object) + } + if js.Properties == nil && js.Type == string(schema.Object) { + // 空参数对象 + } + return &schema.ToolInfo{ + Name: fn.Name, + Desc: fn.Description, + ParamsOneOf: schema.NewParamsOneOfByJSONSchema(&js), + }, nil +} + +type mcpBridgeTool struct { + info *schema.ToolInfo + name string + agent *agent.Agent + holder *ConversationHolder + record ExecutionRecorder + chunk func(toolName, toolCallID, chunk string) + invokeNotify *ToolInvokeNotifyHolder + einoAgentName string +} + +func (m *mcpBridgeTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + _ = ctx + return m.info, nil +} + +func (m *mcpBridgeTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (out string, err error) { + _ = opts + toolCallID := compose.GetToolCallID(ctx) + defer func() { + if m.invokeNotify == nil { + return + } + tid := strings.TrimSpace(toolCallID) + if tid == "" { + return + } + success := err == nil && !strings.HasPrefix(out, ToolErrorPrefix) + body := out + if err != nil { + success = false + } else if strings.HasPrefix(out, ToolErrorPrefix) { + success = false + body = strings.TrimPrefix(out, ToolErrorPrefix) + } + m.invokeNotify.Fire(tid, m.name, m.einoAgentName, success, body, err) + }() + return runMCPToolInvocation(ctx, m.agent, m.holder, m.name, argumentsInJSON, m.record, m.chunk) +} + +// runMCPToolInvocation 与 mcpBridgeTool.InvokableRun 共用。 +func runMCPToolInvocation( + ctx context.Context, + ag *agent.Agent, + holder *ConversationHolder, + toolName string, + argumentsInJSON string, + record ExecutionRecorder, + chunk func(toolName, toolCallID, chunk string), +) (string, error) { + var args map[string]interface{} + if argumentsInJSON != "" && argumentsInJSON != "null" { + if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { + // Return soft error (nil error) so the eino graph continues and the LLM can self-correct, + // instead of a hard error that terminates the iteration loop. + return ToolErrorPrefix + fmt.Sprintf( + "Invalid tool arguments JSON: %s\n\nPlease ensure the arguments are a valid JSON object "+ + "(double-quoted keys, matched braces, no trailing commas) and retry.\n\n"+ + "（工具参数 JSON 解析失败：%s。请确保 arguments 是合法的 JSON 对象并重试。）", + err.Error(), err.Error()), nil + } + } + if args == nil { + args = map[string]interface{}{} + } + + if chunk != nil { + toolCallID := compose.GetToolCallID(ctx) + if toolCallID != "" { + if existing, ok := ctx.Value(security.ToolOutputCallbackCtxKey).(security.ToolOutputCallback); ok && existing != nil { + ctx = context.WithValue(ctx, security.ToolOutputCallbackCtxKey, security.ToolOutputCallback(func(c string) { + existing(c) + if strings.TrimSpace(c) == "" { + return + } + chunk(toolName, toolCallID, c) + })) + } else { + ctx = context.WithValue(ctx, security.ToolOutputCallbackCtxKey, security.ToolOutputCallback(func(c string) { + if strings.TrimSpace(c) == "" { + return + } + chunk(toolName, toolCallID, c) + })) + } + } + } + + res, err := ag.ExecuteMCPToolForConversation(ctx, holder.Get(), toolName, args) + if err != nil { + return "", err + } + if res == nil { + return "", nil + } + if res.ExecutionID != "" && record != nil { + record(res.ExecutionID, compose.GetToolCallID(ctx)) + } + if res.IsError { + return ToolErrorPrefix + res.Result, nil + } + return res.Result, nil +} + +// UnknownToolReminderHandler 供 compose.ToolsNodeConfig.UnknownToolsHandler 使用： +// 模型请求了未注册的工具名时，返回一个「软错误」工具结果（nil error）， +// 让模型在同一轮继续自我修正，避免触发 run-loop 级别的 full rerun。 +// 不进行名称猜测或映射，避免误执行。 +func UnknownToolReminderHandler() func(ctx context.Context, name, input string) (string, error) { + return func(ctx context.Context, name, input string) (string, error) { + _ = ctx + _ = input + requested := strings.TrimSpace(name) + // Return a soft tool-result error so the graph keeps running and the LLM + // can correct tool name/arguments within the same run. + return ToolErrorPrefix + unknownToolReminderText(requested), nil + } +} + +func unknownToolReminderText(requested string) string { + if requested == "" { + requested = "(empty)" + } + return fmt.Sprintf(`The tool name %q is not registered for this agent. + +Please retry using only names that appear in the tool definitions for this turn (exact match, case-sensitive). Do not invent or rename tools; adjust your plan and continue. + +（工具 %q 未注册：请仅使用本回合上下文中给出的工具名称，须完全一致；请勿自行改写或猜测名称，并继续后续步骤。）`, requested, requested) +} diff --git a/internal/einomcp/mcp_tools_test.go b/internal/einomcp/mcp_tools_test.go new file mode 100644 index 00000000..078c8c04 --- /dev/null +++ b/internal/einomcp/mcp_tools_test.go @@ -0,0 +1,16 @@ +package einomcp + +import ( + "strings" + "testing" +) + +func TestUnknownToolReminderText(t *testing.T) { + s := unknownToolReminderText("bad_tool") + if !strings.Contains(s, "bad_tool") { + t.Fatalf("expected requested name in message: %s", s) + } + if strings.Contains(s, "Tools currently available") { + t.Fatal("unified message must not list tool names") + } +} diff --git a/internal/einomcp/tool_invoke_notify.go b/internal/einomcp/tool_invoke_notify.go new file mode 100644 index 00000000..a776a7bc --- /dev/null +++ b/internal/einomcp/tool_invoke_notify.go @@ -0,0 +1,39 @@ +package einomcp + +import "sync" + +// ToolInvokeNotifyHolder 由 Eino run loop 与 MCP/execute 桥共享；Fire 在工具原始返回时触发。 +// UI 的 tool_result 须等 ADK schema.Tool 事件（reduction 后正文），不在此 holder 的回调里推送。 +type ToolInvokeNotifyHolder struct { + mu sync.RWMutex + fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) +} + +// NewToolInvokeNotifyHolder 创建可在 ToolsFromDefinitions 与 run loop 之间共享的 holder。 +func NewToolInvokeNotifyHolder() *ToolInvokeNotifyHolder { + return &ToolInvokeNotifyHolder{} +} + +// Set 由 runEinoADKAgentLoop 在开始消费 iter 之前调用；可多次覆盖（通常仅一次）。 +func (h *ToolInvokeNotifyHolder) Set(fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error)) { + if h == nil { + return + } + h.mu.Lock() + defer h.mu.Unlock() + h.fn = fn +} + +// Fire 由 mcpBridgeTool 在工具调用返回时调用；若尚未 Set 或 toolCallID 为空则忽略。 +func (h *ToolInvokeNotifyHolder) Fire(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) { + if h == nil { + return + } + h.mu.RLock() + fn := h.fn + h.mu.RUnlock() + if fn == nil { + return + } + fn(toolCallID, toolName, einoAgent, success, content, invokeErr) +} diff --git a/internal/einoobserve/attach.go b/internal/einoobserve/attach.go new file mode 100644 index 00000000..9846a2f7 --- /dev/null +++ b/internal/einoobserve/attach.go @@ -0,0 +1,455 @@ +// Package einoobserve attaches CloudWeGo Eino [callbacks.Handler] to ADK Runner contexts for +// structured logging and optional SSE trace events (eino_trace_*). +package einoobserve + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + "github.com/google/uuid" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" + "go.uber.org/zap" +) + +type ctxSpanKey struct{} + +type ctxOtelSpanKey struct{} + +// Params for attaching per-run callback instrumentation. +type Params struct { + Logger *zap.Logger + Progress func(eventType, message string, data interface{}) + ConversationID string + OrchMode string + OrchestratorName string + RunID string +} + +// AttachAgentRunCallbacks returns ctx wrapped with callbacks.InitCallbacks when enabled. +// Safe to call with nil cfg or disabled cfg (returns ctx unchanged). +func AttachAgentRunCallbacks(ctx context.Context, cfg *config.MultiAgentEinoCallbacksConfig, p Params) context.Context { + if ctx == nil { + return ctx + } + if cfg == nil || !cfg.Enabled { + return ctx + } + mode := cfg.EinoCallbacksModeEffective() + if mode == "off" { + return ctx + } + runID := strings.TrimSpace(p.RunID) + if runID == "" { + runID = uuid.New().String() + } + if p.Progress != nil && cfg.ShouldEmitEinoTraceSSE(mode) { + p.Progress("eino_trace_run", "Eino callbacks session", map[string]interface{}{ + "runId": runID, + "conversationId": strings.TrimSpace(p.ConversationID), + "orchestration": strings.TrimSpace(p.OrchMode), + "orchestratorName": strings.TrimSpace(p.OrchestratorName), + "observeMode": mode, + "source": "eino_callbacks", + }) + } + h := &runHandler{ + cfg: *cfg, + mode: mode, + params: p, + runID: runID, + } + b := callbacks.NewHandlerBuilder(). + OnStartFn(h.onStart). + OnEndFn(h.onEnd). + OnErrorFn(h.onError) + if mode == "full" { + b = b.OnStartWithStreamInputFn(h.onStartStreamIn).OnEndWithStreamOutputFn(h.onEndStreamOut) + } + ri := &callbacks.RunInfo{ + Name: "CyberStrikeADKRun", + Type: strings.TrimSpace(p.OrchMode), + Component: components.Component("AgentSession"), + } + return callbacks.InitCallbacks(ctx, ri, b.Build()) +} + +type runHandler struct { + cfg config.MultiAgentEinoCallbacksConfig + mode string + params Params + runID string + + mu sync.Mutex + spanStack []string + seq atomic.Uint64 +} + +func safeRunInfo(info *callbacks.RunInfo) callbacks.RunInfo { + if info == nil { + return callbacks.RunInfo{ + Name: "unknown", + Type: "unknown", + Component: components.Component("unknown"), + } + } + return *info +} + +func (h *runHandler) genSpanID() string { + return fmt.Sprintf("%s-%d", h.runID, h.seq.Add(1)) +} + +func (h *runHandler) popSpan() (id string) { + h.mu.Lock() + defer h.mu.Unlock() + if len(h.spanStack) == 0 { + return "" + } + id = h.spanStack[len(h.spanStack)-1] + h.spanStack = h.spanStack[:len(h.spanStack)-1] + return id +} + +// popMatching removes the given id from the stack top if it matches; otherwise pops until empty or match (rare ordering mismatch). +func (h *runHandler) popMatching(want string) string { + h.mu.Lock() + defer h.mu.Unlock() + if want == "" { + if len(h.spanStack) == 0 { + return "" + } + id := h.spanStack[len(h.spanStack)-1] + h.spanStack = h.spanStack[:len(h.spanStack)-1] + return id + } + for len(h.spanStack) > 0 { + top := h.spanStack[len(h.spanStack)-1] + h.spanStack = h.spanStack[:len(h.spanStack)-1] + if top == want { + return top + } + } + return want +} + +func (h *runHandler) onStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + ri := safeRunInfo(info) + var parentID string + h.mu.Lock() + if len(h.spanStack) > 0 { + parentID = h.spanStack[len(h.spanStack)-1] + } + spanID := h.genSpanID() + h.spanStack = append(h.spanStack, spanID) + h.mu.Unlock() + + inSum := summarizeCallbackInput(input, h.cfg.EinoCallbacksMaxInputSummaryRunes()) + if h.cfg.OtelTracingActive() { + tracer := otel.Tracer("cyberstrike/eino") + spanName := callbackSpanName(info) + var sp trace.Span + ctx, sp = tracer.Start(ctx, spanName, + trace.WithSpanKind(trace.SpanKindInternal), + trace.WithAttributes( + attribute.String("eino.component", string(ri.Component)), + attribute.String("eino.name", ri.Name), + attribute.String("eino.type", ri.Type), + attribute.String("cyberstrike.run_id", h.runID), + attribute.String("cyberstrike.conversation_id", strings.TrimSpace(h.params.ConversationID)), + attribute.String("cyberstrike.orchestration", strings.TrimSpace(h.params.OrchMode)), + ), + ) + if inSum != "" { + sp.SetAttributes(attribute.String("eino.input.summary", truncateForAttr(inSum, 256))) + } + ctx = context.WithValue(ctx, ctxOtelSpanKey{}, sp) + } + if h.params.Logger != nil { + fields := []zap.Field{ + zap.String("runId", h.runID), + zap.String("spanId", spanID), + zap.String("parentSpanId", parentID), + zap.String("component", string(ri.Component)), + zap.String("name", ri.Name), + zap.String("type", ri.Type), + zap.String("phase", "start"), + } + if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil { + if sc := sp.SpanContext(); sc.IsValid() { + fields = append(fields, + zap.String("trace_id", sc.TraceID().String()), + zap.String("otel_span_id", sc.SpanID().String()), + ) + } + } + if h.cfg.ZapVerbose { + h.params.Logger.Debug("eino_callback", append(fields, zap.String("inputSummary", inSum))...) + } else { + h.params.Logger.Info("eino_callback", fields...) + } + } + if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) { + h.params.Progress("eino_trace_start", "", map[string]interface{}{ + "runId": h.runID, + "spanId": spanID, + "parentSpanId": parentID, + "conversationId": strings.TrimSpace(h.params.ConversationID), + "orchestration": strings.TrimSpace(h.params.OrchMode), + "component": string(ri.Component), + "name": ri.Name, + "type": ri.Type, + "ts": time.Now().UTC().Format(time.RFC3339Nano), + "inputSummary": inSum, + "source": "eino_callbacks", + }) + } + ctx = context.WithValue(ctx, ctxSpanKey{}, spanID) + return ctx +} + +func (h *runHandler) onEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + ri := safeRunInfo(info) + spanID, _ := ctx.Value(ctxSpanKey{}).(string) + if spanID == "" { + spanID = h.popSpan() + } else { + spanID = h.popMatching(spanID) + } + outSum := summarizeCallbackOutput(output, h.cfg.EinoCallbacksMaxOutputSummaryRunes()) + if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil { + if outSum != "" { + sp.SetAttributes(attribute.String("eino.output.summary", truncateForAttr(outSum, 256))) + } + sp.SetStatus(codes.Ok, "") + sp.End() + } + if h.params.Logger != nil { + fields := []zap.Field{ + zap.String("runId", h.runID), + zap.String("spanId", spanID), + zap.String("component", string(ri.Component)), + zap.String("name", ri.Name), + zap.String("type", ri.Type), + zap.String("phase", "end"), + } + if h.cfg.ZapVerbose { + h.params.Logger.Debug("eino_callback", append(fields, zap.String("outputSummary", outSum))...) + } else { + h.params.Logger.Info("eino_callback", fields...) + } + } + if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) { + h.params.Progress("eino_trace_end", "", map[string]interface{}{ + "runId": h.runID, + "spanId": spanID, + "conversationId": strings.TrimSpace(h.params.ConversationID), + "orchestration": strings.TrimSpace(h.params.OrchMode), + "component": string(ri.Component), + "name": ri.Name, + "type": ri.Type, + "ts": time.Now().UTC().Format(time.RFC3339Nano), + "outputSummary": outSum, + "source": "eino_callbacks", + }) + } + return ctx +} + +func (h *runHandler) onError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + ri := safeRunInfo(info) + spanID, _ := ctx.Value(ctxSpanKey{}).(string) + if spanID == "" { + spanID = h.popSpan() + } else { + spanID = h.popMatching(spanID) + } + msg := "" + if err != nil { + msg = truncateRunes(err.Error(), h.cfg.EinoCallbacksMaxOutputSummaryRunes()) + } + if sp, ok := ctx.Value(ctxOtelSpanKey{}).(trace.Span); ok && sp != nil { + if err != nil { + sp.RecordError(err) + } + sp.SetStatus(codes.Error, msg) + sp.End() + } + if h.params.Logger != nil { + h.params.Logger.Warn("eino_callback_error", + zap.String("runId", h.runID), + zap.String("spanId", spanID), + zap.String("component", string(ri.Component)), + zap.String("name", ri.Name), + zap.String("type", ri.Type), + zap.Error(err), + ) + } + if h.params.Progress != nil && h.cfg.ShouldEmitEinoTraceSSE(h.mode) { + h.params.Progress("eino_trace_error", msg, map[string]interface{}{ + "runId": h.runID, + "spanId": spanID, + "conversationId": strings.TrimSpace(h.params.ConversationID), + "orchestration": strings.TrimSpace(h.params.OrchMode), + "component": string(ri.Component), + "name": ri.Name, + "type": ri.Type, + "ts": time.Now().UTC().Format(time.RFC3339Nano), + "error": msg, + "source": "eino_callbacks", + }) + } + return ctx +} + +func (h *runHandler) onStartStreamIn(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + ri := safeRunInfo(info) + if input != nil { + input.Close() + } + if h.params.Logger != nil { + h.params.Logger.Debug("eino_callback_stream_in", + zap.String("runId", h.runID), + zap.String("component", string(ri.Component)), + zap.String("name", ri.Name), + ) + } + return ctx +} + +func (h *runHandler) onEndStreamOut(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + ri := safeRunInfo(info) + if output != nil { + output.Close() + } + if h.params.Logger != nil { + h.params.Logger.Debug("eino_callback_stream_out", + zap.String("runId", h.runID), + zap.String("component", string(ri.Component)), + zap.String("name", ri.Name), + ) + } + return ctx +} + +func callbackSpanName(info *callbacks.RunInfo) string { + if info == nil { + return "eino.callback" + } + comp := strings.TrimSpace(string(info.Component)) + name := strings.TrimSpace(info.Name) + typ := strings.TrimSpace(info.Type) + if name != "" && comp != "" { + return comp + "/" + name + } + if typ != "" && comp != "" { + return comp + "[" + typ + "]" + } + if comp != "" { + return comp + } + return "eino.callback" +} + +func truncateForAttr(s string, maxRunes int) string { + return truncateRunes(s, maxRunes) +} + +func summarizeCallbackInput(in callbacks.CallbackInput, maxRunes int) string { + if in == nil { + return "" + } + if ai := adk.ConvAgentCallbackInput(in); ai != nil { + parts := []string{"agent"} + if ai.Input != nil { + parts = append(parts, fmt.Sprintf("messages=%d", len(ai.Input.Messages))) + } + if ai.ResumeInfo != nil { + parts = append(parts, "resume=true") + } + return strings.Join(parts, " ") + } + if mi := model.ConvCallbackInput(in); mi != nil { + return fmt.Sprintf("chatModel messages=%d tools=%d", len(mi.Messages), len(mi.Tools)) + } + if ti := tool.ConvCallbackInput(in); ti != nil { + raw := ti.ArgumentsInJSON + return "tool args=" + truncateRunes(raw, maxRunes) + } + b, err := json.Marshal(in) + if err != nil { + return fmt.Sprintf("%T", in) + } + return truncateRunes(string(b), maxRunes) +} + +func summarizeCallbackOutput(out callbacks.CallbackOutput, maxRunes int) string { + if out == nil { + return "" + } + if ao := adk.ConvAgentCallbackOutput(out); ao != nil { + return "agent_events=stream" + } + if mo := model.ConvCallbackOutput(out); mo != nil && mo.Message != nil { + s := "" + if mo.Message.Content != "" { + s = mo.Message.Content + } + if mo.TokenUsage != nil { + return fmt.Sprintf("tokens total=%d completion=%d prompt=%d text=%s", + mo.TokenUsage.TotalTokens, mo.TokenUsage.CompletionTokens, mo.TokenUsage.PromptTokens, + truncateRunes(s, minInt(120, maxRunes))) + } + return "assistant len=" + itoa(len(s)) + } + if to := tool.ConvCallbackOutput(out); to != nil { + if to.Response != "" { + return truncateRunes(to.Response, maxRunes) + } + if to.ToolOutput != nil { + return "tool_result multimodal" + } + } + b, err := json.Marshal(out) + if err != nil { + return fmt.Sprintf("%T", out) + } + return truncateRunes(string(b), maxRunes) +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} + +func itoa(n int) string { + return fmt.Sprintf("%d", n) +} + +func truncateRunes(s string, maxRunes int) string { + if maxRunes <= 0 { + return "" + } + r := []rune(s) + if len(r) <= maxRunes { + return s + } + return string(r[:maxRunes]) + "…" +} diff --git a/internal/einoobserve/attach_test.go b/internal/einoobserve/attach_test.go new file mode 100644 index 00000000..d12290a2 --- /dev/null +++ b/internal/einoobserve/attach_test.go @@ -0,0 +1,49 @@ +package einoobserve + +import ( + "context" + "testing" + + "cyberstrike-ai/internal/config" +) + +func TestAttachAgentRunCallbacks_Disabled(t *testing.T) { + ctx := context.Background() + cfg := &config.MultiAgentEinoCallbacksConfig{Enabled: false} + out := AttachAgentRunCallbacks(ctx, cfg, Params{}) + if out != ctx { + t.Fatalf("expected same ctx when disabled") + } +} + +func TestAttachAgentRunCallbacksUsesProvidedRunID(t *testing.T) { + emit := true + var gotRunID string + ctx := context.Background() + cfg := &config.MultiAgentEinoCallbacksConfig{Enabled: true, Mode: "sse", SseTraceToClient: &emit} + + AttachAgentRunCallbacks(ctx, cfg, Params{ + RunID: "run-shared", + Progress: func(eventType, _ string, data interface{}) { + if eventType != "eino_trace_run" { + return + } + if m, ok := data.(map[string]interface{}); ok { + gotRunID, _ = m["runId"].(string) + } + }, + }) + + if gotRunID != "run-shared" { + t.Fatalf("runId = %q, want run-shared", gotRunID) + } +} + +func TestTruncateRunes(t *testing.T) { + if got := truncateRunes("abc", 10); got != "abc" { + t.Fatalf("got %q", got) + } + if got := truncateRunes("abcdefghij", 4); got != "abcd…" { + t.Fatalf("got %q", got) + } +} diff --git a/internal/einoobserve/otel.go b/internal/einoobserve/otel.go new file mode 100644 index 00000000..05800abd --- /dev/null +++ b/internal/einoobserve/otel.go @@ -0,0 +1,111 @@ +package einoobserve + +import ( + "context" + "fmt" + "strings" + "sync" + + "cyberstrike-ai/internal/config" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" + "go.opentelemetry.io/otel/exporters/stdout/stdouttrace" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.26.0" + "go.uber.org/zap" +) + +var ( + otelMu sync.Mutex + otelShutdown func(context.Context) error + otelInitialized bool +) + +// InitOtelFromConfig installs the global OpenTelemetry TracerProvider when +// eino_callbacks.otel is enabled and exporter is not none. Safe to call multiple times. +func InitOtelFromConfig(cfg *config.MultiAgentEinoCallbacksConfig, log *zap.Logger) (shutdown func(context.Context) error, err error) { + shutdown = func(context.Context) error { return nil } + if cfg == nil || !cfg.OtelTracingActive() { + return shutdown, nil + } + + otelMu.Lock() + defer otelMu.Unlock() + if otelInitialized { + if otelShutdown != nil { + return otelShutdown, nil + } + return shutdown, nil + } + + oc := cfg.Otel + expKind := oc.OtelExporterEffective() + ctx := context.Background() + + var exporter sdktrace.SpanExporter + switch expKind { + case "stdout": + exporter, err = stdouttrace.New() + if err != nil { + return shutdown, fmt.Errorf("eino otel stdout exporter: %w", err) + } + case "otlphttp": + ep := strings.TrimSpace(oc.OTLPEndpoint) + if ep == "" { + ep = "localhost:4318" + } + exporter, err = otlptracehttp.New(ctx, + otlptracehttp.WithEndpoint(ep), + otlptracehttp.WithURLPath("/v1/traces"), + ) + if err != nil { + return shutdown, fmt.Errorf("eino otel otlphttp exporter: %w", err) + } + default: + return shutdown, nil + } + + res, err := resource.New(ctx, + resource.WithAttributes( + semconv.ServiceName(oc.ServiceNameEffective()), + ), + ) + if err != nil { + return shutdown, fmt.Errorf("eino otel resource: %w", err) + } + + sampler := sdktrace.ParentBased(sdktrace.TraceIDRatioBased(oc.SampleRatioEffective())) + tp := sdktrace.NewTracerProvider( + sdktrace.WithBatcher(exporter), + sdktrace.WithResource(res), + sdktrace.WithSampler(sampler), + ) + otel.SetTracerProvider(tp) + + otelShutdown = tp.Shutdown + otelInitialized = true + if log != nil { + log.Info("eino otel: tracer provider initialized", + zap.String("exporter", expKind), + zap.String("service", oc.ServiceNameEffective()), + zap.Float64("sample_ratio", oc.SampleRatioEffective()), + ) + } + return otelShutdown, nil +} + +// ShutdownOtel flushes and shuts down the global TracerProvider if it was installed. +func ShutdownOtel(ctx context.Context) error { + otelMu.Lock() + fn := otelShutdown + otelShutdown = nil + inited := otelInitialized + otelInitialized = false + otelMu.Unlock() + if !inited || fn == nil { + return nil + } + return fn(ctx) +} diff --git a/internal/logger/logger.go b/internal/logger/logger.go new file mode 100644 index 00000000..7e306fab --- /dev/null +++ b/internal/logger/logger.go @@ -0,0 +1,68 @@ +package logger + +import ( + "os" + + "go.uber.org/zap" + "go.uber.org/zap/zapcore" +) + +type Logger struct { + *zap.Logger +} + +func New(level, output string) *Logger { + var zapLevel zapcore.Level + switch level { + case "debug": + zapLevel = zapcore.DebugLevel + case "info": + zapLevel = zapcore.InfoLevel + case "warn": + zapLevel = zapcore.WarnLevel + case "error": + zapLevel = zapcore.ErrorLevel + default: + zapLevel = zapcore.InfoLevel + } + + config := zap.NewProductionConfig() + config.Level = zap.NewAtomicLevelAt(zapLevel) + config.EncoderConfig.TimeKey = "timestamp" + config.EncoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder + + var writeSyncer zapcore.WriteSyncer + if output == "stdout" { + writeSyncer = zapcore.AddSync(os.Stdout) + } else { + file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666) + if err != nil { + writeSyncer = zapcore.AddSync(os.Stdout) + } else { + writeSyncer = zapcore.AddSync(file) + } + } + + core := zapcore.NewCore( + zapcore.NewJSONEncoder(config.EncoderConfig), + writeSyncer, + zapLevel, + ) + + logger := zap.New(core, zap.AddCaller(), zap.AddStacktrace(zapcore.ErrorLevel)) + + return &Logger{Logger: logger} +} + +func (l *Logger) Fatal(msg string, fields ...interface{}) { + zapFields := make([]zap.Field, 0, len(fields)) + for _, f := range fields { + switch v := f.(type) { + case error: + zapFields = append(zapFields, zap.Error(v)) + default: + zapFields = append(zapFields, zap.Any("field", v)) + } + } + l.Logger.Fatal(msg, zapFields...) +} diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go new file mode 100644 index 00000000..d38bcaae --- /dev/null +++ b/internal/mcp/builtin/constants.go @@ -0,0 +1,195 @@ +package builtin + +// 内置工具名称常量 +// 所有代码中使用内置工具名称的地方都应该使用这些常量，而不是硬编码字符串 +const ( + // 漏洞管理工具 + ToolRecordVulnerability = "record_vulnerability" + ToolListVulnerabilities = "list_vulnerabilities" + ToolGetVulnerability = "get_vulnerability" + + // 资产管理工具 + ToolCreateAsset = "create_asset" + ToolGetAsset = "get_asset" + ToolQueryAssets = "query_assets" + ToolUpdateAsset = "update_asset" + ToolDeleteAsset = "delete_asset" + ToolCompleteAssetScan = "complete_asset_scan" + + // 项目黑板（事实）工具 + ToolUpsertProjectFact = "upsert_project_fact" + ToolGetProjectFact = "get_project_fact" + ToolListProjectFacts = "list_project_facts" + ToolSearchProjectFacts = "search_project_facts" + ToolDeprecateProjectFact = "deprecate_project_fact" + ToolRestoreProjectFact = "restore_project_fact" + + // 知识库工具 + ToolListKnowledgeRiskTypes = "list_knowledge_risk_types" + ToolSearchKnowledgeBase = "search_knowledge_base" + + // 视觉分析（本地图片 → VL 模型 → 文本摘要） + ToolAnalyzeImage = "analyze_image" + + // 长耗时工具执行控制（后台 execution 查询/等待/取消） + ToolGetToolExecution = "get_tool_execution" + ToolWaitToolExecution = "wait_tool_execution" + ToolCancelToolExecution = "cancel_tool_execution" + + // WebShell 助手工具（AI 在 WebShell 管理 - AI 助手 中使用） + ToolWebshellExec = "webshell_exec" + ToolWebshellFileList = "webshell_file_list" + ToolWebshellFileRead = "webshell_file_read" + ToolWebshellFileWrite = "webshell_file_write" + + // WebShell 连接管理工具（用于通过 MCP 管理 webshell 连接） + ToolManageWebshellList = "manage_webshell_list" + ToolManageWebshellAdd = "manage_webshell_add" + ToolManageWebshellUpdate = "manage_webshell_update" + ToolManageWebshellDelete = "manage_webshell_delete" + ToolManageWebshellTest = "manage_webshell_test" + + // 批量任务队列（与 Web 端批量任务一致，供模型创建/启停/查询队列） + ToolBatchTaskList = "batch_task_list" + ToolBatchTaskGet = "batch_task_get" + ToolBatchTaskCreate = "batch_task_create" + ToolBatchTaskStart = "batch_task_start" + ToolBatchTaskRerun = "batch_task_rerun" + ToolBatchTaskPause = "batch_task_pause" + ToolBatchTaskDelete = "batch_task_delete" + ToolBatchTaskUpdateMetadata = "batch_task_update_metadata" + ToolBatchTaskUpdateSchedule = "batch_task_update_schedule" + ToolBatchTaskScheduleEnabled = "batch_task_schedule_enabled" + ToolBatchTaskAdd = "batch_task_add_task" + ToolBatchTaskUpdate = "batch_task_update_task" + ToolBatchTaskRemove = "batch_task_remove_task" + + // C2 工具集（合并同类项，8 个统一工具） + ToolC2Listener = "c2_listener" // 监听器管理（create/start/stop/list/get/update/delete） + ToolC2Session = "c2_session" // 会话管理（list/get/set_sleep/kill/delete） + ToolC2Task = "c2_task" // 任务下发（统一 task_type 参数） + ToolC2TaskManage = "c2_task_manage" // 任务管理（get_result/wait/list/cancel） + ToolC2Payload = "c2_payload" // Payload 生成（oneliner/build） + ToolC2Event = "c2_event" // 事件查询 + ToolC2Profile = "c2_profile" // Malleable Profile 管理（list/get/create/update/delete） + ToolC2File = "c2_file" // 文件管理（list/get_result） +) + +// IsBuiltinTool 检查工具名称是否是内置工具 +func IsBuiltinTool(toolName string) bool { + switch toolName { + case ToolRecordVulnerability, + ToolListVulnerabilities, + ToolGetVulnerability, + ToolCreateAsset, + ToolGetAsset, + ToolQueryAssets, + ToolUpdateAsset, + ToolDeleteAsset, + ToolCompleteAssetScan, + ToolUpsertProjectFact, + ToolGetProjectFact, + ToolListProjectFacts, + ToolSearchProjectFacts, + ToolDeprecateProjectFact, + ToolRestoreProjectFact, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase, + ToolAnalyzeImage, + ToolGetToolExecution, + ToolWaitToolExecution, + ToolCancelToolExecution, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite, + ToolManageWebshellList, + ToolManageWebshellAdd, + ToolManageWebshellUpdate, + ToolManageWebshellDelete, + ToolManageWebshellTest, + ToolBatchTaskList, + ToolBatchTaskGet, + ToolBatchTaskCreate, + ToolBatchTaskStart, + ToolBatchTaskRerun, + ToolBatchTaskPause, + ToolBatchTaskDelete, + ToolBatchTaskUpdateMetadata, + ToolBatchTaskUpdateSchedule, + ToolBatchTaskScheduleEnabled, + ToolBatchTaskAdd, + ToolBatchTaskUpdate, + ToolBatchTaskRemove, + // C2 工具 + ToolC2Listener, + ToolC2Session, + ToolC2Task, + ToolC2TaskManage, + ToolC2Payload, + ToolC2Event, + ToolC2Profile, + ToolC2File: + return true + default: + return false + } +} + +// GetAllBuiltinTools 返回所有内置工具名称列表 +func GetAllBuiltinTools() []string { + return []string{ + ToolRecordVulnerability, + ToolListVulnerabilities, + ToolGetVulnerability, + ToolCreateAsset, + ToolGetAsset, + ToolQueryAssets, + ToolUpdateAsset, + ToolDeleteAsset, + ToolCompleteAssetScan, + ToolUpsertProjectFact, + ToolGetProjectFact, + ToolListProjectFacts, + ToolSearchProjectFacts, + ToolDeprecateProjectFact, + ToolRestoreProjectFact, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase, + ToolAnalyzeImage, + ToolGetToolExecution, + ToolWaitToolExecution, + ToolCancelToolExecution, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite, + ToolManageWebshellList, + ToolManageWebshellAdd, + ToolManageWebshellUpdate, + ToolManageWebshellDelete, + ToolManageWebshellTest, + ToolBatchTaskList, + ToolBatchTaskGet, + ToolBatchTaskCreate, + ToolBatchTaskStart, + ToolBatchTaskRerun, + ToolBatchTaskPause, + ToolBatchTaskDelete, + ToolBatchTaskUpdateMetadata, + ToolBatchTaskUpdateSchedule, + ToolBatchTaskScheduleEnabled, + ToolBatchTaskAdd, + ToolBatchTaskUpdate, + ToolBatchTaskRemove, + // C2 工具 + ToolC2Listener, + ToolC2Session, + ToolC2Task, + ToolC2TaskManage, + ToolC2Payload, + ToolC2Event, + ToolC2Profile, + ToolC2File, + } +} diff --git a/internal/mcp/client_sdk.go b/internal/mcp/client_sdk.go new file mode 100644 index 00000000..0d7ebfb3 --- /dev/null +++ b/internal/mcp/client_sdk.go @@ -0,0 +1,475 @@ +// Package mcp 外部 MCP 客户端 - 基于官方 go-sdk 实现，保证协议兼容性 +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "os/exec" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "go.uber.org/zap" +) + +const ( + clientName = "CyberStrikeAI" + clientVersion = "1.0.0" +) + +// sdkClient 基于官方 MCP Go SDK 的外部 MCP 客户端，实现 ExternalMCPClient 接口 +type sdkClient struct { + session *mcp.ClientSession + client *mcp.Client + logger *zap.Logger + mu sync.RWMutex + status string // "disconnected", "connecting", "connected", "error" +} + +// newSDKClientFromSession 用已连接成功的 session 构造（供 createSDKClient 内部使用） +func newSDKClientFromSession(session *mcp.ClientSession, client *mcp.Client, logger *zap.Logger) *sdkClient { + return &sdkClient{ + session: session, + client: client, + logger: logger, + status: "connected", + } +} + +// lazySDKClient 延迟连接：Initialize() 时才调用官方 SDK 建立连接，对外实现 ExternalMCPClient +type lazySDKClient struct { + serverCfg config.ExternalMCPServerConfig + logger *zap.Logger + sessionCancel context.CancelFunc + inner ExternalMCPClient // connected SDK client + mu sync.RWMutex + status string +} + +func newLazySDKClient(serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) *lazySDKClient { + return &lazySDKClient{ + serverCfg: serverCfg, + logger: logger, + status: "connecting", + } +} + +func (c *lazySDKClient) setStatus(s string) { + c.mu.Lock() + defer c.mu.Unlock() + c.status = s +} + +func (c *lazySDKClient) GetStatus() string { + c.mu.RLock() + defer c.mu.RUnlock() + if c.inner != nil { + return c.inner.GetStatus() + } + return c.status +} + +func (c *lazySDKClient) IsConnected() bool { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner != nil { + return inner.IsConnected() + } + return false +} + +func (c *lazySDKClient) Initialize(ctx context.Context) error { + c.mu.Lock() + if c.inner != nil { + c.mu.Unlock() + return nil + } + c.mu.Unlock() + + sessionCtx, sessionCancel := context.WithCancel(context.Background()) + type connectResult struct { + inner ExternalMCPClient + err error + } + resultCh := make(chan connectResult) + abandoned := make(chan struct{}) + go func() { + inner, err := createSDKClient(sessionCtx, c.serverCfg, c.logger) + select { + case resultCh <- connectResult{inner: inner, err: err}: + case <-abandoned: + if inner != nil { + _ = inner.Close() + } + sessionCancel() + } + }() + + var result connectResult + select { + case result = <-resultCh: + case <-ctx.Done(): + close(abandoned) + sessionCancel() + c.setStatus("error") + return ctx.Err() + } + + if err := ctx.Err(); err != nil { + sessionCancel() + if result.inner != nil { + _ = result.inner.Close() + } + c.setStatus("error") + return err + } + + if result.err != nil { + sessionCancel() + c.setStatus("error") + return result.err + } + + c.mu.Lock() + if c.inner != nil { + c.mu.Unlock() + sessionCancel() + if result.inner != nil { + _ = result.inner.Close() + } + return nil + } + c.inner = result.inner + c.sessionCancel = sessionCancel + c.mu.Unlock() + c.setStatus("connected") + return nil +} + +func (c *lazySDKClient) ListTools(ctx context.Context) ([]Tool, error) { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner == nil { + return nil, fmt.Errorf("未连接") + } + return inner.ListTools(ctx) +} + +func (c *lazySDKClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner == nil { + return nil, fmt.Errorf("未连接") + } + return inner.CallTool(ctx, name, args) +} + +func (c *lazySDKClient) Close() error { + c.mu.Lock() + inner := c.inner + sessionCancel := c.sessionCancel + c.inner = nil + c.sessionCancel = nil + c.mu.Unlock() + c.setStatus("disconnected") + if sessionCancel != nil { + sessionCancel() + } + if inner != nil { + return inner.Close() + } + return nil +} + +// markDisconnected 在检测到传输层断连时关闭底层 session，避免 IsConnected 仍返回 true。 +func (c *lazySDKClient) markDisconnected() { + c.mu.Lock() + inner := c.inner + sessionCancel := c.sessionCancel + c.inner = nil + c.sessionCancel = nil + c.mu.Unlock() + if sessionCancel != nil { + sessionCancel() + } + if inner != nil { + _ = inner.Close() + } + c.setStatus("disconnected") +} + +func (c *sdkClient) setStatus(s string) { + c.mu.Lock() + defer c.mu.Unlock() + c.status = s +} + +func (c *sdkClient) GetStatus() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.status +} + +func (c *sdkClient) IsConnected() bool { + return c.GetStatus() == "connected" +} + +func (c *sdkClient) Initialize(ctx context.Context) error { + // sdkClient 由 createSDKClient 在 Connect 成功后才创建，因此 Initialize 时已经连接 + // 此方法仅用于满足 ExternalMCPClient 接口，实际连接在 createSDKClient 中完成 + return nil +} + +func (c *sdkClient) ListTools(ctx context.Context) ([]Tool, error) { + if c.session == nil { + return nil, fmt.Errorf("未连接") + } + res, err := c.session.ListTools(ctx, nil) + if err != nil { + return nil, err + } + if res == nil { + return nil, nil + } + return sdkToolsToOur(res.Tools), nil +} + +func (c *sdkClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + if c.session == nil { + return nil, fmt.Errorf("未连接") + } + params := &mcp.CallToolParams{ + Name: name, + Arguments: args, + } + res, err := c.session.CallTool(ctx, params) + if err != nil { + return nil, err + } + return sdkCallToolResultToOurs(res), nil +} + +func (c *sdkClient) Close() error { + c.setStatus("disconnected") + if c.session != nil { + err := c.session.Close() + c.session = nil + return err + } + return nil +} + +// sdkToolsToOur 将 SDK 的 []*mcp.Tool 转为我们的 []Tool +func sdkToolsToOur(tools []*mcp.Tool) []Tool { + if len(tools) == 0 { + return nil + } + out := make([]Tool, 0, len(tools)) + for _, t := range tools { + if t == nil { + continue + } + schema := make(map[string]interface{}) + if t.InputSchema != nil { + // SDK InputSchema 可能为 *jsonschema.Schema 或 map，统一转为 map + if m, ok := t.InputSchema.(map[string]interface{}); ok { + schema = m + } else { + _ = json.Unmarshal(mustJSON(t.InputSchema), &schema) + } + } + desc := t.Description + shortDesc := desc + if t.Annotations != nil && t.Annotations.Title != "" { + shortDesc = t.Annotations.Title + } + out = append(out, Tool{ + Name: t.Name, + Description: desc, + ShortDescription: shortDesc, + InputSchema: schema, + }) + } + return out +} + +// sdkCallToolResultToOurs 将 SDK 的 *mcp.CallToolResult 转为我们的 *ToolResult +func sdkCallToolResultToOurs(res *mcp.CallToolResult) *ToolResult { + if res == nil { + return &ToolResult{Content: []Content{}} + } + content := sdkContentToOurs(res.Content) + return &ToolResult{ + Content: content, + IsError: res.IsError, + } +} + +func sdkContentToOurs(list []mcp.Content) []Content { + if len(list) == 0 { + return nil + } + out := make([]Content, 0, len(list)) + for _, c := range list { + switch v := c.(type) { + case *mcp.TextContent: + out = append(out, Content{Type: "text", Text: v.Text}) + default: + out = append(out, Content{Type: "text", Text: fmt.Sprintf("%v", c)}) + } + } + return out +} + +func mustJSON(v interface{}) []byte { + b, _ := json.Marshal(v) + return b +} + +// createSDKClient 根据配置创建并连接外部 MCP 客户端（使用官方 SDK），返回实现 ExternalMCPClient 的 *sdkClient +// 若连接失败返回 (nil, error)。ctx 用于连接超时与取消。 +func createSDKClient(ctx context.Context, serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) (ExternalMCPClient, error) { + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + transport := serverCfg.GetTransportType() + if transport == "" { + return nil, fmt.Errorf("配置缺少 command 或 url，且未指定 type/transport") + } + + // 构造 ClientOptions：KeepAlive 心跳 + var clientOpts *mcp.ClientOptions + if serverCfg.KeepAlive > 0 { + clientOpts = &mcp.ClientOptions{ + KeepAlive: time.Duration(serverCfg.KeepAlive) * time.Second, + } + } + + client := mcp.NewClient(&mcp.Implementation{ + Name: clientName, + Version: clientVersion, + }, clientOpts) + + var t mcp.Transport + switch transport { + case "stdio": + if serverCfg.Command == "" { + return nil, fmt.Errorf("stdio 模式需要配置 command") + } + // 必须用 exec.Command 而非 CommandContext：doConnect 返回后 ctx 会被 cancel， + // 若用 CommandContext(ctx) 会立刻杀掉子进程，导致 ListTools 等后续请求失败、显示 0 工具 + cmd := exec.Command(serverCfg.Command, serverCfg.Args...) + if len(serverCfg.Env) > 0 { + cmd.Env = append(cmd.Env, envMapToSlice(serverCfg.Env)...) + } + ct := &mcp.CommandTransport{Command: cmd} + if serverCfg.TerminateDuration > 0 { + ct.TerminateDuration = time.Duration(serverCfg.TerminateDuration) * time.Second + } + t = ct + case "sse": + if serverCfg.URL == "" { + return nil, fmt.Errorf("sse 模式需要配置 url") + } + // SSE 是长连接（GET 流持续打开），不能设置 http.Client.Timeout（会在超时后杀掉整个连接导致 EOF）。 + // 超时由每次 ListTools/CallTool 的 context 单独控制。 + httpClient := httpClientForLongLived(serverCfg.Headers) + t = &mcp.SSEClientTransport{ + Endpoint: serverCfg.URL, + HTTPClient: httpClient, + } + case "http": + if serverCfg.URL == "" { + return nil, fmt.Errorf("http 模式需要配置 url") + } + httpClient := httpClientWithTimeoutAndHeaders(timeout, serverCfg.Headers) + st := &mcp.StreamableClientTransport{ + Endpoint: serverCfg.URL, + HTTPClient: httpClient, + } + if serverCfg.MaxRetries > 0 { + st.MaxRetries = serverCfg.MaxRetries + } + t = st + default: + return nil, fmt.Errorf("不支持的传输模式: %s（支持: stdio, sse, http）", transport) + } + + session, err := client.Connect(ctx, t, nil) + if err != nil { + return nil, fmt.Errorf("连接失败: %w", err) + } + + return newSDKClientFromSession(session, client, logger), nil +} + +func envMapToSlice(env map[string]string) []string { + m := make(map[string]string) + for _, s := range os.Environ() { + if i := strings.IndexByte(s, '='); i > 0 { + m[s[:i]] = s[i+1:] + } + } + for k, v := range env { + m[k] = v + } + out := make([]string, 0, len(m)) + for k, v := range m { + out = append(out, k+"="+v) + } + return out +} + +func httpClientWithTimeoutAndHeaders(timeout time.Duration, headers map[string]string) *http.Client { + transport := http.DefaultTransport + if len(headers) > 0 { + transport = &headerRoundTripper{ + headers: headers, + base: http.DefaultTransport, + } + } + return &http.Client{ + Timeout: timeout, + Transport: transport, + } +} + +// httpClientForLongLived 创建不设超时的 HTTP 客户端，用于 SSE 等长连接传输。 +// SSE 的 GET 流会持续打开，http.Client.Timeout 会在超时后强制关闭连接导致 EOF。 +// 超时由调用方通过 context 控制。 +func httpClientForLongLived(headers map[string]string) *http.Client { + transport := http.DefaultTransport + if len(headers) > 0 { + transport = &headerRoundTripper{ + headers: headers, + base: http.DefaultTransport, + } + } + return &http.Client{ + Transport: transport, + // 不设 Timeout，SSE 长连接的超时由 per-request context 控制 + } +} + +type headerRoundTripper struct { + headers map[string]string + base http.RoundTripper +} + +func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + for k, v := range h.headers { + req.Header.Set(k, v) + } + return h.base.RoundTrip(req) +} diff --git a/internal/mcp/connection_recovery.go b/internal/mcp/connection_recovery.go new file mode 100644 index 00000000..a2ed9bfb --- /dev/null +++ b/internal/mcp/connection_recovery.go @@ -0,0 +1,192 @@ +package mcp + +import ( + "context" + "errors" + "io" + "strings" + "time" + + "go.uber.org/zap" +) + +const ( + // externalReconnectMinInterval 两次自动重连之间的最短间隔 + externalReconnectMinInterval = 30 * time.Second + // externalReconnectMaxBackoff 指数退避上限 + externalReconnectMaxBackoff = 5 * time.Minute +) + +// isConnectionDeadError 判断错误是否表示底层传输已断开（而非调用方主动取消或超时）。 +func isConnectionDeadError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + if errors.Is(err, io.EOF) { + return true + } + s := strings.ToLower(err.Error()) + return strings.Contains(s, "eof") || + strings.Contains(s, "client is closing") || + strings.Contains(s, "connection closed") || + strings.Contains(s, "connection reset") || + strings.Contains(s, "broken pipe") +} + +// handleConnectionDead 在 ListTools/CallTool 等操作失败且判定为断连时，标记客户端并调度重连。 +func (m *ExternalMCPManager) handleConnectionDead(name string, client ExternalMCPClient, err error) { + if !isConnectionDeadError(err) { + return + } + m.logger.Warn("检测到外部MCP连接已断开，将尝试自动重连", + zap.String("name", name), + zap.Error(err), + ) + m.markClientDisconnected(name, client, err) + m.scheduleReconnect(name) +} + +func (m *ExternalMCPManager) markClientDisconnected(name string, client ExternalMCPClient, err error) { + if lazy, ok := client.(*lazySDKClient); ok { + lazy.markDisconnected() + } + m.mu.Lock() + if err != nil { + m.errors[name] = "连接已断开: " + err.Error() + } + m.mu.Unlock() + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() +} + +func (m *ExternalMCPManager) onClientConnected(name string) { + m.clearReconnectState(name) +} + +func (m *ExternalMCPManager) clearReconnectState(name string) { + m.reconnectMu.Lock() + delete(m.reconnectAttempts, name) + delete(m.reconnectLastTry, name) + delete(m.reconnecting, name) + m.reconnectMu.Unlock() +} + +func (m *ExternalMCPManager) reconnectBackoff(attempts int) time.Duration { + if attempts <= 0 { + return 0 + } + d := externalReconnectMinInterval + for i := 1; i < attempts && d < externalReconnectMaxBackoff; i++ { + d *= 2 + } + if d > externalReconnectMaxBackoff { + return externalReconnectMaxBackoff + } + return d +} + +func (m *ExternalMCPManager) scheduleReconnect(name string) { + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + m.mu.RUnlock() + if !enabled { + return + } + go m.tryReconnect(name) +} + +func (m *ExternalMCPManager) tryReconnect(name string) { + m.reconnectMu.Lock() + if m.reconnecting[name] { + m.reconnectMu.Unlock() + return + } + attempts := m.reconnectAttempts[name] + if wait := m.reconnectBackoff(attempts); wait > 0 { + if last, ok := m.reconnectLastTry[name]; ok { + if elapsed := time.Since(last); elapsed < wait { + remaining := wait - elapsed + m.reconnectMu.Unlock() + m.scheduleReconnectAfter(name, remaining) + return + } + } + } + m.reconnecting[name] = true + m.reconnectMu.Unlock() + + defer func() { + m.reconnectMu.Lock() + delete(m.reconnecting, name) + m.reconnectMu.Unlock() + }() + + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + client, hasClient := m.clients[name] + connecting := hasClient && client.GetStatus() == "connecting" + m.mu.RUnlock() + + if !enabled { + m.logger.Debug("跳过自动重连（外部MCP已停用）", zap.String("name", name)) + return + } + if connecting { + m.logger.Debug("跳过自动重连（连接正在进行中）", zap.String("name", name)) + return + } + + m.reconnectMu.Lock() + m.reconnectLastTry[name] = time.Now() + m.reconnectAttempts[name] = attempts + 1 + attemptNum := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + + m.logger.Info("正在自动重连外部MCP", + zap.String("name", name), + zap.Int("attempt", attemptNum), + ) + + if err := m.startClient(name, true); err != nil { + m.logger.Warn("自动重连外部MCP失败", + zap.String("name", name), + zap.Error(err), + ) + } +} + +// scheduleReconnectAfterFailure 在自动重连失败后，按当前退避间隔预约下一次重试。 +func (m *ExternalMCPManager) scheduleReconnectAfterFailure(name string) { + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + m.mu.RUnlock() + if !enabled { + return + } + m.reconnectMu.Lock() + wait := m.reconnectBackoff(m.reconnectAttempts[name]) + m.reconnectMu.Unlock() + m.logger.Info("自动重连失败，将按退避间隔再次尝试", + zap.String("name", name), + zap.Duration("after", wait), + ) + m.scheduleReconnectAfter(name, wait) +} + +// scheduleReconnectAfter 在 delay 后触发 tryReconnect（delay<=0 时立即执行）。 +func (m *ExternalMCPManager) scheduleReconnectAfter(name string, delay time.Duration) { + if delay <= 0 { + go m.tryReconnect(name) + return + } + time.AfterFunc(delay, func() { + m.tryReconnect(name) + }) +} diff --git a/internal/mcp/connection_recovery_test.go b/internal/mcp/connection_recovery_test.go new file mode 100644 index 00000000..f04e4622 --- /dev/null +++ b/internal/mcp/connection_recovery_test.go @@ -0,0 +1,215 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "io" + "testing" + "time" + + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func TestIsConnectionDeadError(t *testing.T) { + t.Parallel() + cases := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"eof", io.EOF, true}, + {"wrapped eof", fmt.Errorf("connection closed: %w", io.EOF), true}, + {"client closing", errors.New(`calling "tools/list": client is closing: EOF`), true}, + {"connection reset", errors.New("read tcp: connection reset by peer"), true}, + {"canceled", context.Canceled, false}, + {"deadline", context.DeadlineExceeded, false}, + {"other", errors.New("invalid params"), false}, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := isConnectionDeadError(tc.err); got != tc.want { + t.Fatalf("isConnectionDeadError(%v) = %v, want %v", tc.err, got, tc.want) + } + }) + } +} + +func TestLazySDKClient_MarkDisconnected(t *testing.T) { + c := &lazySDKClient{status: "connected"} + c.inner = &sdkClient{status: "connected"} + c.markDisconnected() + if c.IsConnected() { + t.Fatal("expected disconnected after markDisconnected") + } + if c.GetStatus() != "disconnected" { + t.Fatalf("expected status disconnected, got %s", c.GetStatus()) + } +} + +func TestHandleConnectionDead_MarksLazyClientDisconnected(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "dead-mcp" + cfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: true, + } + m.mu.Lock() + m.configs[name] = cfg + client := newLazySDKClient(cfg, logger) + client.inner = &sdkClient{status: "connected"} + client.status = "connected" + m.clients[name] = client + m.mu.Unlock() + + deadErr := errors.New(`connection closed: calling "tools/list": client is closing: EOF`) + m.handleConnectionDead(name, client, deadErr) + + if client.IsConnected() { + t.Fatal("expected disconnected after handleConnectionDead") + } + if m.GetError(name) == "" { + t.Fatal("expected error message to be recorded") + } + counts := m.GetToolCounts() + if counts[name] != 0 { + t.Fatalf("expected tool count 0 after disconnect, got %d", counts[name]) + } +} + +func TestReconnectBackoff(t *testing.T) { + t.Parallel() + if d := (&ExternalMCPManager{}).reconnectBackoff(0); d != 0 { + t.Fatalf("attempt 0: got %v", d) + } + if d := (&ExternalMCPManager{}).reconnectBackoff(1); d != externalReconnectMinInterval { + t.Fatalf("attempt 1: got %v", d) + } + if d := (&ExternalMCPManager{}).reconnectBackoff(10); d != externalReconnectMaxBackoff { + t.Fatalf("attempt 10: got %v, want cap %v", d, externalReconnectMaxBackoff) + } +} + +func TestTryReconnect_RateLimited(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "rate-limited" + m.reconnectMu.Lock() + m.reconnectLastTry[name] = time.Now() + m.reconnectAttempts[name] = 2 + m.reconnectMu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 2 { + t.Fatalf("rate limited reconnect should not increment attempts, got %d", attempts) + } +} + +func TestTryReconnect_SkipsWhenDisabled(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "disabled-mcp" + m.mu.Lock() + m.configs[name] = config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: false, + } + m.mu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 0 { + t.Fatalf("disabled MCP should not increment reconnect attempts, got %d", attempts) + } +} + +func TestTryReconnect_SkipsWhenConnecting(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "connecting-mcp" + cfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: true, + } + client := newLazySDKClient(cfg, logger) + client.setStatus("connecting") + + m.mu.Lock() + m.configs[name] = cfg + m.clients[name] = client + m.mu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 0 { + t.Fatalf("connecting MCP should not increment reconnect attempts, got %d", attempts) + } +} + +func TestStartClientAutoReconnect_SkipsWhenDisabled(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + m.stopRefresh = make(chan struct{}) + + name := "stopped" + m.mu.Lock() + m.configs[name] = config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: false, + } + m.mu.Unlock() + + if err := m.startClient(name, true); err != nil { + t.Fatalf("startClient: %v", err) + } + + m.mu.RLock() + cfg := m.configs[name] + _, hasClient := m.clients[name] + m.mu.RUnlock() + if cfg.ExternalMCPEnable { + t.Fatal("auto reconnect should not enable stopped MCP") + } + if hasClient { + t.Fatal("auto reconnect should not create client when disabled") + } +} + +func TestOnClientConnected_ClearsReconnectState(t *testing.T) { + m := &ExternalMCPManager{ + reconnectAttempts: map[string]int{"x": 3}, + reconnectLastTry: map[string]time.Time{"x": time.Now()}, + reconnecting: map[string]bool{"x": true}, + } + m.onClientConnected("x") + + m.reconnectMu.Lock() + defer m.reconnectMu.Unlock() + if len(m.reconnectAttempts) != 0 || len(m.reconnectLastTry) != 0 || len(m.reconnecting) != 0 { + t.Fatal("expected reconnect state cleared") + } +} diff --git a/internal/mcp/execution_control_tools.go b/internal/mcp/execution_control_tools.go new file mode 100644 index 00000000..7d74af25 --- /dev/null +++ b/internal/mcp/execution_control_tools.go @@ -0,0 +1,296 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/mcp/builtin" +) + +const ( + defaultExecutionWaitTimeout = 60 * time.Second + maxExecutionWaitTimeout = 10 * time.Minute + defaultPartialPreviewBytes = 4096 + maxPartialPreviewBytes = 64 * 1024 +) + +// RegisterExecutionControlTools exposes execution handle operations to Eino as +// ordinary MCP tools. This keeps the agent loop native: the model calls a tool, +// receives a bounded result, and may call wait_tool_execution again if needed. +func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) { + if server == nil { + return + } + + server.RegisterTool(Tool{ + Name: builtin.ToolGetToolExecution, + Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后，凭 execution_id 继续查看进度。", + ShortDescription: "查询后台工具执行状态", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + exec := lookupToolExecution(server, external, id) + if exec == nil { + return textToolResult("未找到该 execution_id: "+id, true), nil + } + return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil + }) + + server.RegisterTool(Tool{ + Name: builtin.ToolWaitToolExecution, + Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限；若仍未完成，会返回当前状态，模型可稍后再次调用。", + ShortDescription: "有界等待后台工具执行", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数，默认 60，最大 600"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout) + snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait) + if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) { + return textToolResult("等待 execution 失败: "+err.Error(), true), nil + } + if snap == nil || snap.Execution == nil { + return textToolResult("未找到该 execution_id: "+id, true), nil + } + body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args)) + if errors.Is(err, ErrExecutionWaitTimeout) { + body += "\n\n本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。" + } + return textToolResult(body, false), nil + }) + + server.RegisterTool(Tool{ + Name: builtin.ToolCancelToolExecution, + Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。", + ShortDescription: "取消后台工具执行", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "reason": map[string]interface{}{"type": "string", "description": "取消原因，可选，会写入终止说明"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + reason := stringArg(args, "reason") + if server.CancelToolExecutionWithNote(id, reason) { + return textToolResult("已请求取消内部工具 execution: "+id, false), nil + } + if external != nil && external.CancelToolExecutionWithNote(id, reason) { + return textToolResult("已请求取消外部 MCP execution: "+id, false), nil + } + return textToolResult("未找到进行中的 execution，或该 execution 已结束: "+id, true), nil + }) +} + +func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) { + if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil { + return server.executionService.Wait(ctx, id, wait) + } + if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil { + return external.executionService.Wait(ctx, id, wait) + } + if server != nil && server.executionService != nil { + if snap, err := server.executionService.Get(id); err == nil { + return snap, nil + } + } + if external != nil && external.executionService != nil { + return external.executionService.Get(id) + } + exec := lookupToolExecution(server, external, id) + if exec == nil { + return nil, fmt.Errorf("execution not found: %s", id) + } + return &ExecutionSnapshot{Execution: exec}, nil +} + +func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution { + if server != nil { + if exec, ok := server.GetExecution(id); ok && exec != nil { + return exec + } + } + if external != nil { + if exec, ok := external.GetExecution(id); ok && exec != nil { + return exec + } + } + return nil +} + +type executionFormatOptions struct { + includePartialOutput bool + partialMaxBytes int +} + +func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions { + includePartial := true + if raw, ok := args["include_partial_output"]; ok { + if b, ok := raw.(bool); ok { + includePartial = b + } else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" { + includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes") + } + } + maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes) + return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes} +} + +func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string { + if exec == nil { + return "execution: null" + } + payload := map[string]interface{}{ + "execution_id": exec.ID, + "tool": exec.ToolName, + "status": exec.Status, + "started_at": exec.StartTime.Format(time.RFC3339), + } + if exec.EndTime != nil { + payload["ended_at"] = exec.EndTime.Format(time.RFC3339) + } + if exec.Duration > 0 { + payload["duration"] = exec.Duration.String() + } + if exec.Error != "" { + payload["error"] = exec.Error + } + if exec.Result != nil { + payload["result"] = ToolResultPlainText(exec.Result) + payload["is_error"] = exec.Result.IsError + } + if opts.includePartialOutput && exec.PartialOutput != "" { + partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes) + payload["partial_output"] = partial + payload["partial_output_bytes"] = exec.PartialOutputBytes + payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput)) + if exec.PartialOutputUpdatedAt != nil { + payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339) + } + } + b, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error) + } + return string(b) +} + +func tailStringBytes(s string, maxBytes int) string { + if maxBytes <= 0 { + maxBytes = defaultPartialPreviewBytes + } + b := []byte(s) + if len(b) <= maxBytes { + return s + } + return string(b[len(b)-maxBytes:]) +} + +func textToolResult(text string, isErr bool) *ToolResult { + return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr} +} + +func stringArg(args map[string]interface{}, key string) string { + if args == nil { + return "" + } + raw, ok := args[key] + if !ok || raw == nil { + return "" + } + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + default: + return strings.TrimSpace(fmt.Sprint(v)) + } +} + +func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration { + if args == nil { + return def + } + var seconds float64 + switch v := args[key].(type) { + case int: + seconds = float64(v) + case int64: + seconds = float64(v) + case float64: + seconds = v + case json.Number: + f, _ := v.Float64() + seconds = f + case string: + f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64) + seconds = f + } + if seconds <= 0 { + return def + } + d := time.Duration(seconds * float64(time.Second)) + if max > 0 && d > max { + return max + } + return d +} + +func intArg(args map[string]interface{}, key string, def, max int) int { + if args == nil { + return def + } + var n int + switch v := args[key].(type) { + case int: + n = v + case int64: + n = int(v) + case float64: + n = int(v) + case json.Number: + i, _ := v.Int64() + n = int(i) + case string: + i, _ := strconv.Atoi(strings.TrimSpace(v)) + n = i + } + if n <= 0 { + return def + } + if max > 0 && n > max { + return max + } + return n +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go new file mode 100644 index 00000000..1e2ad6d5 --- /dev/null +++ b/internal/mcp/execution_service.go @@ -0,0 +1,625 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/authctx" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +const ( + ToolExecutionStatusQueued = "queued" + ToolExecutionStatusRunning = "running" + ToolExecutionStatusCompleted = "completed" + ToolExecutionStatusFailed = "failed" + ToolExecutionStatusCancelled = "cancelled" + ToolExecutionStatusHardTimeout = "hard_timeout" + ToolExecutionStatusOrphaned = "orphaned" +) + +var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout") + +// ExecutionRunFunc is the blocking operation owned by a worker. +type ExecutionRunFunc func(context.Context) (*ToolResult, error) + +type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error) + +// ExecutionDoneFunc observes the final persisted state. It is invoked once, +// including for late completions after an agent has stopped waiting. +type ExecutionDoneFunc func(*ToolExecution) + +type ExecutionRequest struct { + ID string + ToolName string + Arguments map[string]interface{} + ConversationID string + OwnerUserID string + HardTimeout time.Duration + PreRun ExecutionPreRunFunc + Run ExecutionRunFunc + OnDone ExecutionDoneFunc +} + +type ExecutionHandle struct { + ID string +} + +type ExecutionSnapshot struct { + Execution *ToolExecution +} + +type executionEntry struct { + exec *ToolExecution + cancel context.CancelFunc + done chan struct{} + preRun ExecutionPreRunFunc + run ExecutionRunFunc + result *ToolResult + err error +} + +// ExecutionService keeps Eino-facing tool calls synchronous while moving the +// untrusted blocking work into cancellable workers with explicit execution IDs. +type ExecutionService struct { + storage MonitorStorage + logger *zap.Logger + + mu sync.Mutex + entries map[string]*executionEntry + abortUserNotes map[string]string + maxInMemory int + resultMaxBytes int + spillRootDir string +} + +func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService { + if logger == nil { + logger = zap.NewNop() + } + return &ExecutionService{ + storage: storage, + logger: logger, + entries: make(map[string]*executionEntry), + abortUserNotes: make(map[string]string), + maxInMemory: 1000, + resultMaxBytes: DefaultToolResultMaxBytes, + } +} + +func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.resultMaxBytes = maxBytes +} + +// ConfigureToolResultSpillRoot sets the reduction-compatible root used when +// oversized tool results are spilled to local files (empty → tmp/reduction). +func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.spillRootDir = strings.TrimSpace(rootDir) +} + +func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) { + if s == nil { + return nil, fmt.Errorf("execution service is nil") + } + if req.Run == nil { + return nil, fmt.Errorf("execution run func is nil") + } + id := strings.TrimSpace(req.ID) + if id == "" { + id = uuid.New().String() + } + start := time.Now() + exec := &ToolExecution{ + ID: id, + ToolName: strings.TrimSpace(req.ToolName), + Arguments: cloneArgsMap(req.Arguments), + Status: ToolExecutionStatusQueued, + StartTime: start, + ConversationID: strings.TrimSpace(req.ConversationID), + OwnerUserID: strings.TrimSpace(req.OwnerUserID), + } + if exec.ConversationID == "" { + exec.ConversationID = MCPConversationIDFromContext(ctx) + } + if exec.OwnerUserID == "" { + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + exec.OwnerUserID = principal.UserID + } + } + + runCtx := detachedExecutionContext(ctx) + var cancel context.CancelFunc + if req.HardTimeout > 0 { + runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout) + } else { + runCtx, cancel = context.WithCancel(runCtx) + } + entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run} + + s.mu.Lock() + if _, exists := s.entries[id]; exists { + s.mu.Unlock() + cancel() + return nil, fmt.Errorf("execution already exists: %s", id) + } + s.entries[id] = entry + s.cleanupOldEntriesLocked() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(exec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id)) + } + } + notifyToolRunBegin(ctx, id) + + go s.runWorker(runCtx, entry, req.OnDone) + return &ExecutionHandle{ID: id}, nil +} + +func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) { + id := entry.exec.ID + ctx = WithMCPExecutionID(ctx, id) + if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" { + ctx = WithMCPConversationID(ctx, conv) + } + var release func() + defer func() { + if release != nil { + release() + } + entry.cancel() + notifyToolRunEnd(ctx, id) + close(entry.done) + }() + + if entry.preRun != nil { + var preErr error + release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec)) + if preErr != nil { + s.finishEntry(ctx, entry, nil, preErr, onDone) + return + } + } + s.markEntryRunning(entry) + + result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) { + return nilSafeRun(ctx, entry) + }) + s.finishEntry(ctx, entry, result, err, onDone) +} + +func (s *ExecutionService) markEntryRunning(entry *executionEntry) { + if s == nil || entry == nil || entry.exec == nil { + return + } + s.mu.Lock() + if !isExecutionTerminal(entry.exec.Status) { + entry.exec.Status = ToolExecutionStatusRunning + } + runningExec := cloneToolExecution(entry.exec) + s.mu.Unlock() + if s.storage != nil { + if err := s.storage.SaveToolExecution(runningExec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID)) + } + } +} + +func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) { + id := entry.exec.ID + cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err) + + now := time.Now() + s.mu.Lock() + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ConversationID: entry.exec.ConversationID, + ExecutionID: id, + } + if ctx != nil { + if pid := MCPProjectIDFromContext(ctx); pid != "" { + spill.ProjectID = pid + } + if conv := MCPConversationIDFromContext(ctx); conv != "" { + spill.ConversationID = conv + } + } + result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill) + entry.result = result + entry.err = err + entry.exec.EndTime = &now + entry.exec.Duration = now.Sub(entry.exec.StartTime) + if err != nil { + switch { + case errors.Is(err, context.DeadlineExceeded): + entry.exec.Status = ToolExecutionStatusHardTimeout + entry.exec.Error = "工具执行超过硬超时限制" + case errors.Is(err, context.Canceled): + entry.exec.Status = ToolExecutionStatusCancelled + entry.exec.Error = "已手动终止或任务已取消" + default: + entry.exec.Status = ToolExecutionStatusFailed + entry.exec.Error = err.Error() + } + } else if result != nil && result.IsError { + if cancelledWithUserNote { + entry.exec.Status = ToolExecutionStatusCancelled + entry.exec.Error = "" + } else if isBackgroundWaitToolResult(result) { + entry.exec.Status = ToolExecutionStatusCompleted + entry.exec.Error = "" + } else { + entry.exec.Status = ToolExecutionStatusFailed + entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果") + } + entry.exec.Result = result + } else { + entry.exec.Status = ToolExecutionStatusCompleted + if result == nil { + result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}} + entry.result = result + } + entry.exec.Result = result + } + finalExec := cloneToolExecution(entry.exec) + s.mu.Unlock() + + if s.storage != nil { + if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id)) + } + } + if onDone != nil { + onDone(finalExec) + } +} + +func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) { + if entry == nil { + return nil, fmt.Errorf("execution entry is nil") + } + if entry.run == nil { + return nil, fmt.Errorf("execution run func not wired") + } + return entry.run(ctx) +} + +func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) { + defer func() { + if r := recover(); r != nil { + if logger != nil { + logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack")) + } + err = fmt.Errorf("tool execution panic: %v", r) + } + }() + return fn() +} + +func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) { + entry := s.getEntry(executionID) + if entry == nil { + return s.getPersistedSnapshot(executionID) + } + if isExecutionTerminal(entry.exec.Status) { + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + } + + var timeoutCh <-chan time.Time + var timer *time.Timer + if timeout > 0 { + timer = time.NewTimer(timeout) + timeoutCh = timer.C + defer timer.Stop() + } + + select { + case <-entry.done: + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + case <-timeoutCh: + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout + case <-ctxDone(ctx): + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err() + } +} + +func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) { + entry := s.getEntry(executionID) + if entry != nil { + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + } + return s.getPersistedSnapshot(executionID) +} + +func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool { + id := strings.TrimSpace(executionID) + if s == nil || id == "" || chunk == "" { + return false + } + now := time.Now() + s.mu.Lock() + defer s.mu.Unlock() + entry := s.entries[id] + if entry == nil || entry.exec == nil { + return false + } + appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now) + return true +} + +func (s *ExecutionService) Cancel(executionID, note string) bool { + id := strings.TrimSpace(executionID) + if id == "" || s == nil { + return false + } + s.mu.Lock() + entry := s.entries[id] + if entry == nil || isExecutionTerminal(entry.exec.Status) { + s.mu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + s.abortUserNotes[id] = strings.TrimSpace(note) + } + cancel := entry.cancel + s.mu.Unlock() + if cancel != nil { + cancel() + } + return true +} + +func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} { + if s == nil { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + out := make(map[string]struct{}) + for id, entry := range s.entries { + if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) { + out[id] = struct{}{} + } + } + if len(out) == 0 { + return nil + } + return out +} + +func (s *ExecutionService) CancelAll(note string) { + if s == nil { + return + } + s.mu.Lock() + cancels := make([]context.CancelFunc, 0, len(s.entries)) + for id, entry := range s.entries { + if entry == nil || isExecutionTerminal(entry.exec.Status) { + continue + } + if strings.TrimSpace(note) != "" { + s.abortUserNotes[id] = strings.TrimSpace(note) + } + if entry.cancel != nil { + cancels = append(cancels, entry.cancel) + } + } + s.mu.Unlock() + for _, cancel := range cancels { + cancel() + } +} + +func (s *ExecutionService) getEntry(executionID string) *executionEntry { + if s == nil { + return nil + } + id := strings.TrimSpace(executionID) + if id == "" { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + return s.entries[id] +} + +func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) { + id := strings.TrimSpace(executionID) + if id == "" { + return nil, fmt.Errorf("execution_id is required") + } + if s != nil && s.storage != nil { + exec, err := s.storage.GetToolExecution(id) + if err == nil && exec != nil { + return &ExecutionSnapshot{Execution: exec}, nil + } + if err != nil { + return nil, err + } + } + return nil, fmt.Errorf("execution not found: %s", id) +} + +func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(s.takeAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + if err != nil { + *err = nil + } + if result != nil { + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + } + return true +} + +func (s *ExecutionService) takeAbortUserNote(id string) string { + s.mu.Lock() + defer s.mu.Unlock() + note := s.abortUserNotes[id] + delete(s.abortUserNotes, id) + return note +} + +func (s *ExecutionService) cleanupOldEntriesLocked() { + if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory { + return + } + type oldEntry struct { + id string + startTime time.Time + } + var terminal []oldEntry + for id, entry := range s.entries { + if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) { + terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime}) + } + } + for len(s.entries) > s.maxInMemory && len(terminal) > 0 { + oldest := 0 + for i := 1; i < len(terminal); i++ { + if terminal[i].startTime.Before(terminal[oldest].startTime) { + oldest = i + } + } + delete(s.entries, terminal[oldest].id) + terminal = append(terminal[:oldest], terminal[oldest+1:]...) + } +} + +func firstToolResultText(result *ToolResult, fallback string) string { + if result != nil { + for _, c := range result.Content { + if strings.TrimSpace(c.Text) != "" { + return c.Text + } + } + } + return fallback +} + +func isBackgroundWaitToolResult(result *ToolResult) bool { + text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result))) + if text == "" { + return false + } + hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`) + hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") || + strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) || + strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`) + hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") || + strings.Contains(text, "本次等待已到达") || + strings.Contains(text, "wait_timeout:") || + strings.Contains(text, "background execution") || + strings.Contains(text, "still running") || + strings.Contains(text, "仍未完成") + return hasExecutionID && hasRunningStatus && hasSoftWaitSignal +} + +func isExecutionTerminal(status string) bool { + switch strings.TrimSpace(strings.ToLower(status)) { + case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned: + return true + default: + return false + } +} + +func ctxDone(ctx context.Context) <-chan struct{} { + if ctx == nil { + return nil + } + return ctx.Done() +} + +func detachedExecutionContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return context.WithoutCancel(ctx) +} + +func cloneArgsMap(in map[string]interface{}) map[string]interface{} { + if in == nil { + return map[string]interface{}{} + } + out := make(map[string]interface{}, len(in)) + for k, v := range in { + out[k] = v + } + return out +} + +func cloneToolExecution(in *ToolExecution) *ToolExecution { + if in == nil { + return nil + } + out := *in + out.Arguments = cloneArgsMap(in.Arguments) + if in.Result != nil { + res := *in.Result + if in.Result.Content != nil { + res.Content = append([]Content(nil), in.Result.Content...) + } + out.Result = &res + } + if in.EndTime != nil { + t := *in.EndTime + out.EndTime = &t + } + if in.PartialOutputUpdatedAt != nil { + t := *in.PartialOutputUpdatedAt + out.PartialOutputUpdatedAt = &t + } + return &out +} + +func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) { + if exec == nil || chunk == "" { + return + } + if maxBytes <= 0 { + maxBytes = defaultPartialOutputMaxBytes + } + exec.PartialOutputBytes += int64(len([]byte(chunk))) + combined := exec.PartialOutput + chunk + if len([]byte(combined)) > maxBytes { + b := []byte(combined) + combined = string(b[len(b)-maxBytes:]) + exec.PartialOutputTruncated = true + } + exec.PartialOutput = combined + t := updatedAt + exec.PartialOutputUpdatedAt = &t +} diff --git a/internal/mcp/execution_service_test.go b/internal/mcp/execution_service_test.go new file mode 100644 index 00000000..29f981c5 --- /dev/null +++ b/internal/mcp/execution_service_test.go @@ -0,0 +1,41 @@ +package mcp + +import ( + "context" + "testing" +) + +func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) { + service := NewExecutionService(nil, nil) + handle, err := service.Submit(context.Background(), ExecutionRequest{ + ToolName: "wait_tool_execution", + Run: func(context.Context) (*ToolResult, error) { + return &ToolResult{ + Content: []Content{{Type: "text", Text: `{ + "execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7", + "tool": "exec", + "status": "running" +} + +本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。`}}, + IsError: true, + }, nil + }, + }) + if err != nil { + t.Fatalf("Submit: %v", err) + } + snap, err := service.Wait(context.Background(), handle.ID, 0) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if snap == nil || snap.Execution == nil { + t.Fatal("missing execution snapshot") + } + if snap.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted) + } + if snap.Execution.Result == nil || !snap.Execution.Result.IsError { + t.Fatal("model-facing result should remain IsError") + } +} diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go new file mode 100644 index 00000000..f51b0457 --- /dev/null +++ b/internal/mcp/external_manager.go @@ -0,0 +1,1615 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +const ( + // externalToolListCacheTTL 已连接外部 MCP 的工具列表缓存有效期，避免每次 API 请求都打远程 ListTools。 + externalToolListCacheTTL = 60 * time.Second + // externalToolCountRefreshInterval 后台刷新工具数量的间隔（仅刷新缓存过期或缺失的客户端）。 + externalToolCountRefreshInterval = 60 * time.Second +) + +// toolListCacheEntry 外部 MCP 工具列表缓存条目 +type toolListCacheEntry struct { + tools []Tool + updatedAt time.Time +} + +// listToolsInflight 合并同一 MCP 上并发的 ListTools 请求 +type listToolsInflight struct { + done chan struct{} + tools []Tool + err error +} + +type ExternalMCPResilienceConfig struct { + MaxConcurrentPerServer int + MaxConcurrentTotal int + CircuitFailureThreshold int + CircuitCooldown time.Duration +} + +type externalMCPServerRuntime struct { + semaphore chan struct{} + consecutiveFailures int + circuitOpenUntil time.Time +} + +// ExternalMCPManager 外部MCP管理器 +type ExternalMCPManager struct { + clients map[string]ExternalMCPClient + configs map[string]config.ExternalMCPServerConfig + logger *zap.Logger + storage MonitorStorage // 可选的持久化存储 + executions map[string]*ToolExecution // 执行记录 + stats map[string]*ToolStats // 工具统计信息 + errors map[string]string // 错误信息 + toolCounts map[string]int // 工具数量缓存 + toolCountsMu sync.RWMutex // 工具数量缓存的锁 + toolCache map[string]toolListCacheEntry // 工具列表缓存：MCP名称 -> 工具列表 + toolCacheMu sync.RWMutex // 工具列表缓存的锁 + listToolsMu sync.Mutex + listToolsInflight map[string]*listToolsInflight + stopRefresh chan struct{} // 停止后台刷新的信号 + refreshWg sync.WaitGroup // 等待后台刷新goroutine完成 + refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积 + mu sync.RWMutex + runningCancels map[string]context.CancelFunc + abortUserNotes map[string]string + reconnectMu sync.Mutex + reconnecting map[string]bool + reconnectLastTry map[string]time.Time + reconnectAttempts map[string]int + toolAuthorizer func(context.Context, string, map[string]interface{}) error + executionService *ExecutionService + toolWaitTimeout time.Duration + toolResultMaxBytes int + spillRootDir string + resilience ExternalMCPResilienceConfig + serverRuntimes map[string]*externalMCPServerRuntime + globalSemaphore chan struct{} +} + +// NewExternalMCPManager 创建外部MCP管理器 +func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager { + return NewExternalMCPManagerWithStorage(logger, nil) +} + +// SetToolAuthorizer installs the policy decision point for all external MCP +// invocations. App wiring configures this before any Agent can call a tool. +func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + m.mu.Lock() + m.toolAuthorizer = authorizer + m.mu.Unlock() +} + +// NewExternalMCPManagerWithStorage 创建外部MCP管理器（带持久化存储） +func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager { + manager := &ExternalMCPManager{ + clients: make(map[string]ExternalMCPClient), + configs: make(map[string]config.ExternalMCPServerConfig), + logger: logger, + storage: storage, + executions: make(map[string]*ToolExecution), + stats: make(map[string]*ToolStats), + errors: make(map[string]string), + toolCounts: make(map[string]int), + toolCache: make(map[string]toolListCacheEntry), + listToolsInflight: make(map[string]*listToolsInflight), + stopRefresh: make(chan struct{}), + runningCancels: make(map[string]context.CancelFunc), + abortUserNotes: make(map[string]string), + reconnecting: make(map[string]bool), + reconnectLastTry: make(map[string]time.Time), + reconnectAttempts: make(map[string]int), + toolWaitTimeout: 60 * time.Second, + toolResultMaxBytes: DefaultToolResultMaxBytes, + resilience: ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 2, + MaxConcurrentTotal: 16, + CircuitFailureThreshold: 3, + CircuitCooldown: 60 * time.Second, + }, + serverRuntimes: make(map[string]*externalMCPServerRuntime), + globalSemaphore: make(chan struct{}, 16), + } + manager.executionService = NewExecutionService(storage, logger) + // 启动后台刷新工具数量的goroutine + manager.startToolCountRefresh() + return manager +} + +func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) { + if m == nil { + return + } + m.mu.Lock() + m.toolResultMaxBytes = maxBytes + m.mu.Unlock() + if m.executionService != nil { + m.executionService.ConfigureToolResultMaxBytes(maxBytes) + } +} + +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (m *ExternalMCPManager) ConfigureToolResultSpillRoot(rootDir string) { + if m == nil { + return + } + m.mu.Lock() + m.spillRootDir = strings.TrimSpace(rootDir) + m.mu.Unlock() + if m.executionService != nil { + m.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + +// ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call +// waits for an external MCP execution before returning an execution_id that can +// be polled with wait_tool_execution. seconds<=0 waits until completion. +func (m *ExternalMCPManager) ConfigureToolWaitTimeoutSeconds(seconds int) { + if m == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if seconds <= 0 { + m.toolWaitTimeout = 0 + return + } + m.toolWaitTimeout = time.Duration(seconds) * time.Second +} + +func (m *ExternalMCPManager) ConfigureResilience(cfg ExternalMCPResilienceConfig) { + if m == nil { + return + } + normalized := normalizeExternalMCPResilienceConfig(cfg) + m.mu.Lock() + defer m.mu.Unlock() + m.resilience = normalized + m.serverRuntimes = make(map[string]*externalMCPServerRuntime) + if normalized.MaxConcurrentTotal > 0 { + m.globalSemaphore = make(chan struct{}, normalized.MaxConcurrentTotal) + } else { + m.globalSemaphore = nil + } +} + +func normalizeExternalMCPResilienceConfig(cfg ExternalMCPResilienceConfig) ExternalMCPResilienceConfig { + if cfg.MaxConcurrentPerServer == 0 { + cfg.MaxConcurrentPerServer = 2 + } + if cfg.MaxConcurrentTotal == 0 { + cfg.MaxConcurrentTotal = 16 + } + if cfg.CircuitFailureThreshold == 0 { + cfg.CircuitFailureThreshold = 3 + } + if cfg.CircuitCooldown <= 0 { + cfg.CircuitCooldown = 60 * time.Second + } + if cfg.MaxConcurrentPerServer < 0 { + cfg.MaxConcurrentPerServer = 0 + } + if cfg.MaxConcurrentTotal < 0 { + cfg.MaxConcurrentTotal = 0 + } + return cfg +} + +// LoadConfigs 加载配置 +func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) { + m.mu.Lock() + defer m.mu.Unlock() + + if cfg == nil || cfg.Servers == nil { + return + } + + m.configs = make(map[string]config.ExternalMCPServerConfig) + for name, serverCfg := range cfg.Servers { + m.configs[name] = serverCfg + } +} + +// GetConfigs 获取所有配置 +func (m *ExternalMCPManager) GetConfigs() map[string]config.ExternalMCPServerConfig { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]config.ExternalMCPServerConfig) + for k, v := range m.configs { + result[k] = v + } + return result +} + +// AddOrUpdateConfig 添加或更新配置 +func (m *ExternalMCPManager) AddOrUpdateConfig(name string, serverCfg config.ExternalMCPServerConfig) error { + m.mu.Lock() + defer m.mu.Unlock() + + // 如果已存在客户端，先关闭 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + m.configs[name] = serverCfg + + // 如果启用，自动连接 + if m.isEnabled(serverCfg) { + go m.connectClient(name, serverCfg) + } + + return nil +} + +// RemoveConfig 移除配置 +func (m *ExternalMCPManager) RemoveConfig(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + // 关闭客户端 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + delete(m.configs, name) + m.clearReconnectState(name) + + // 清理工具数量缓存 + m.toolCountsMu.Lock() + delete(m.toolCounts, name) + m.toolCountsMu.Unlock() + + // 清理工具列表缓存 + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() + + return nil +} + +// StartClient 启动客户端（用户手动启动；连接失败不自动重试） +func (m *ExternalMCPManager) StartClient(name string) error { + return m.startClient(name, false) +} + +// startClient 启动客户端。autoReconnect 为 true 时用于断连自愈：尊重停用状态，失败后按退避继续重试。 +func (m *ExternalMCPManager) startClient(name string, autoReconnect bool) error { + m.mu.Lock() + serverCfg, exists := m.configs[name] + m.mu.Unlock() + + if !exists { + return fmt.Errorf("配置不存在: %s", name) + } + + if autoReconnect && !m.isEnabled(serverCfg) { + return nil + } + + // 检查是否已经有连接的客户端 + m.mu.RLock() + existingClient, hasClient := m.clients[name] + m.mu.RUnlock() + + if hasClient { + // 检查客户端是否已连接 + if existingClient.IsConnected() { + // 客户端已连接，直接返回成功（目标状态已达成） + if !autoReconnect { + m.mu.Lock() + serverCfg.ExternalMCPEnable = true + m.configs[name] = serverCfg + m.mu.Unlock() + } + return nil + } + // 如果有客户端但未连接，先关闭 + existingClient.Close() + m.mu.Lock() + delete(m.clients, name) + m.mu.Unlock() + } + + if autoReconnect { + m.mu.RLock() + serverCfg, exists = m.configs[name] + enabled := exists && m.isEnabled(serverCfg) + m.mu.RUnlock() + if !enabled { + return nil + } + } + + // 更新配置为启用 + m.mu.Lock() + serverCfg.ExternalMCPEnable = true + m.configs[name] = serverCfg + // 清除之前的错误信息（重新启动时） + delete(m.errors, name) + m.mu.Unlock() + + // 立即创建客户端并设置为"connecting"状态，这样前端可以立即看到状态 + client := m.createClient(serverCfg) + if client == nil { + return fmt.Errorf("无法创建客户端：不支持的传输模式") + } + + // 设置状态为connecting + m.setClientStatus(client, "connecting") + + // 立即保存客户端，这样前端查询时就能看到"connecting"状态 + m.mu.Lock() + m.clients[name] = client + m.mu.Unlock() + + // 在后台异步进行实际连接 + go func(reconnect bool) { + if err := m.doConnect(name, serverCfg, client); err != nil { + m.logger.Error("连接外部MCP客户端失败", + zap.String("name", name), + zap.Bool("auto_reconnect", reconnect), + zap.Error(err), + ) + // 连接失败，设置状态为error并保存错误信息 + m.setClientStatus(client, "error") + m.mu.Lock() + m.errors[name] = err.Error() + m.mu.Unlock() + // 触发工具数量刷新（连接失败，工具数量应为0） + m.triggerToolCountRefresh() + if reconnect { + m.scheduleReconnectAfterFailure(name) + } + } else { + // 连接成功，清除错误信息 + m.mu.Lock() + delete(m.errors, name) + m.mu.Unlock() + m.onClientConnected(name) + // 异步拉取工具列表（singleflight 去重，结果同时写入 toolCache 与 toolCounts） + go m.refreshToolCache(name, client) + } + }(autoReconnect) + + return nil +} + +// StopClient 停止客户端 +func (m *ExternalMCPManager) StopClient(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + serverCfg, exists := m.configs[name] + if !exists { + return fmt.Errorf("配置不存在: %s", name) + } + + // 关闭客户端 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + // 清除错误信息 + delete(m.errors, name) + + // 更新工具数量缓存（停止后工具数量为0） + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() + + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() + + // 更新配置为禁用 + serverCfg.ExternalMCPEnable = false + m.configs[name] = serverCfg + + m.clearReconnectState(name) + + return nil +} + +// GetClient 获取客户端 +func (m *ExternalMCPManager) GetClient(name string) (ExternalMCPClient, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + + client, exists := m.clients[name] + return client, exists +} + +// GetError 获取错误信息 +func (m *ExternalMCPManager) GetError(name string) string { + m.mu.RLock() + defer m.mu.RUnlock() + + return m.errors[name] +} + +// GetAllTools 获取所有外部MCP的工具 +// 优先从已连接的客户端获取，如果连接断开则返回缓存的工具列表 +// 策略： +// - error 状态：不使用缓存，直接跳过（配置错误或服务不可用） +// - disconnected/connecting 状态：使用缓存（临时断开） +// - connected 状态：正常获取，失败时降级使用缓存 +func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) { + m.mu.RLock() + clients := make(map[string]ExternalMCPClient) + for k, v := range m.clients { + clients[k] = v + } + m.mu.RUnlock() + + var allTools []Tool + var hasError bool + var lastError error + + // 使用较短的超时时间进行快速检查（3秒），避免阻塞 + quickCtx, quickCancel := context.WithTimeout(ctx, 3*time.Second) + defer quickCancel() + + for name, client := range clients { + tools, err := m.getToolsForClient(name, client, quickCtx) + if err != nil { + // 记录错误，但继续处理其他客户端 + hasError = true + if lastError == nil { + lastError = err + } + continue + } + + // 为工具添加前缀，避免冲突 + for _, tool := range tools { + tool.Name = fmt.Sprintf("%s::%s", name, tool.Name) + allTools = append(allTools, tool) + } + } + + // 如果有错误但至少返回了一些工具，不返回错误（部分成功） + if hasError && len(allTools) == 0 { + return nil, fmt.Errorf("获取外部MCP工具失败: %w", lastError) + } + + return allTools, nil +} + +// getToolsForClient 获取指定客户端的工具列表 +// 返回工具列表和错误（如果完全无法获取） +func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPClient, ctx context.Context) ([]Tool, error) { + status := client.GetStatus() + + // error 状态：不使用缓存，直接返回错误 + if status == "error" { + m.logger.Debug("跳过连接失败的外部MCP（不使用缓存）", + zap.String("name", name), + zap.String("status", status), + ) + return nil, fmt.Errorf("外部MCP连接失败: %s", name) + } + + // 已连接：缓存优先，仅在缺失或过期时打远程 ListTools + if client.IsConnected() { + if tools, ok := m.getFreshCachedTools(name); ok { + return tools, nil + } + if tools, ok := m.getAnyCachedTools(name); ok { + m.triggerToolListRefresh(name, client) + return tools, nil + } + tools, err := m.listToolsDeduped(ctx, name, client) + if err != nil { + return m.getCachedTools(name, "连接正常但获取失败", err) + } + return tools, nil + } + + // 未连接：根据状态决定是否使用缓存 + if status == "disconnected" || status == "connecting" { + return m.getCachedTools(name, fmt.Sprintf("客户端临时断开（状态: %s）", status), nil) + } + + // 其他未知状态，不使用缓存 + m.logger.Debug("跳过外部MCP（未知状态）", + zap.String("name", name), + zap.String("status", status), + ) + return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status) +} + +// getCachedTools 获取缓存的工具列表（含空列表缓存） +func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) { + if tools, ok := m.getAnyCachedTools(name); ok { + m.logger.Debug("使用缓存的工具列表", + zap.String("name", name), + zap.String("reason", reason), + zap.Int("count", len(tools)), + zap.Error(originalErr), + ) + return tools, nil + } + + if originalErr != nil { + return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr) + } + return nil, fmt.Errorf("外部MCP无缓存工具: %s", name) +} + +func (m *ExternalMCPManager) isToolCacheFresh(updatedAt time.Time) bool { + return !updatedAt.IsZero() && time.Since(updatedAt) < externalToolListCacheTTL +} + +func cloneTools(tools []Tool) []Tool { + if len(tools) == 0 { + return nil + } + out := make([]Tool, len(tools)) + copy(out, tools) + return out +} + +func (m *ExternalMCPManager) getFreshCachedTools(name string) ([]Tool, bool) { + m.toolCacheMu.RLock() + entry, ok := m.toolCache[name] + m.toolCacheMu.RUnlock() + if !ok || !m.isToolCacheFresh(entry.updatedAt) { + return nil, false + } + return cloneTools(entry.tools), true +} + +func (m *ExternalMCPManager) getAnyCachedTools(name string) ([]Tool, bool) { + m.toolCacheMu.RLock() + entry, ok := m.toolCache[name] + m.toolCacheMu.RUnlock() + if !ok { + return nil, false + } + return cloneTools(entry.tools), true +} + +// listToolsDeduped 对同一 MCP 合并并发 ListTools，并更新 toolCache / toolCounts。 +func (m *ExternalMCPManager) listToolsDeduped(ctx context.Context, name string, client ExternalMCPClient) ([]Tool, error) { + m.listToolsMu.Lock() + if inflight, exists := m.listToolsInflight[name]; exists { + m.listToolsMu.Unlock() + select { + case <-inflight.done: + if inflight.err != nil { + return nil, inflight.err + } + return cloneTools(inflight.tools), nil + case <-ctx.Done(): + return nil, ctx.Err() + } + } + inflight := &listToolsInflight{done: make(chan struct{})} + m.listToolsInflight[name] = inflight + m.listToolsMu.Unlock() + + inflight.tools, inflight.err = client.ListTools(ctx) + if inflight.err == nil { + m.updateToolCache(name, inflight.tools) + } + + m.listToolsMu.Lock() + delete(m.listToolsInflight, name) + close(inflight.done) + m.listToolsMu.Unlock() + + if inflight.err != nil { + m.handleConnectionDead(name, client, inflight.err) + return nil, inflight.err + } + return cloneTools(inflight.tools), nil +} + +// InvalidateToolCache 清除指定外部 MCP 的工具列表缓存（手动刷新时使用） +func (m *ExternalMCPManager) InvalidateToolCache(name string) { + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() +} + +// InvalidateAllToolCaches 清除所有外部 MCP 工具列表缓存 +func (m *ExternalMCPManager) InvalidateAllToolCaches() { + m.toolCacheMu.Lock() + m.toolCache = make(map[string]toolListCacheEntry) + m.toolCacheMu.Unlock() +} + +func (m *ExternalMCPManager) triggerToolListRefresh(name string, client ExternalMCPClient) { + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + _, _ = m.listToolsDeduped(ctx, name, client) + }() +} + +// updateToolCache 更新工具列表缓存与工具数量 +func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) { + stored := cloneTools(tools) + m.toolCacheMu.Lock() + m.toolCache[name] = toolListCacheEntry{tools: stored, updatedAt: time.Now()} + m.toolCacheMu.Unlock() + + m.toolCountsMu.Lock() + m.toolCounts[name] = len(stored) + m.toolCountsMu.Unlock() + + if len(stored) == 0 { + m.logger.Warn("外部MCP返回空工具列表", + zap.String("name", name), + zap.String("hint", "服务可能暂时不可用，工具列表为空"), + ) + } else { + m.logger.Debug("工具列表缓存已更新", + zap.String("name", name), + zap.Int("count", len(stored)), + ) + } +} + +// CallTool 调用外部MCP工具（返回执行ID） +func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { + if m.executionService == nil { + m.executionService = NewExecutionService(m.storage, m.logger) + m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes) + m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir) + } + var ownerUserID string + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + ownerUserID = principal.UserID + } + var mcpName, actualToolName string + var client ExternalMCPClient + handle, err := m.executionService.Submit(ctx, ExecutionRequest{ + ToolName: toolName, + Arguments: args, + ConversationID: MCPConversationIDFromContext(ctx), + OwnerUserID: ownerUserID, + PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + m.mu.RLock() + authorizer := m.toolAuthorizer + m.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("external tool authorization denied: %w", err) + } + } else if authenticated { + return nil, fmt.Errorf("external tool authorization policy is not configured") + } + + // 解析工具名称：name::toolName + if idx := findSubstring(toolName, "::"); idx > 0 { + mcpName = toolName[:idx] + actualToolName = toolName[idx+2:] + } else { + return nil, fmt.Errorf("无效的工具名称格式: %s", toolName) + } + + var exists bool + client, exists = m.GetClient(mcpName) + if !exists { + return nil, fmt.Errorf("外部MCP客户端不存在: %s", mcpName) + } + if err := m.checkExternalMCPCircuit(mcpName); err != nil { + return nil, err + } + + // 检查连接状态，如果未连接或状态为error，不允许调用 + if !client.IsConnected() { + status := client.GetStatus() + if status == "error" { + // 获取错误信息（如果有） + errorMsg := m.GetError(mcpName) + if errorMsg != "" { + return nil, fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg) + } + return nil, fmt.Errorf("外部MCP连接失败: %s", mcpName) + } + return nil, fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status) + } + + release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName) + if acquireErr != nil { + return nil, acquireErr + } + return release, nil + }, + Run: func(runCtx context.Context) (*ToolResult, error) { + result, callErr := client.CallTool(runCtx, actualToolName, args) + if callErr != nil { + m.handleConnectionDead(mcpName, client, callErr) + } + return result, callErr + }, + OnDone: func(exec *ToolExecution) { + failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled + if mcpName != "" { + m.recordExternalMCPResult(mcpName, failed) + } + m.updateStats(toolName, failed) + }, + }) + if err != nil { + return nil, "", err + } + + m.mu.RLock() + waitTimeout := m.toolWaitTimeout + m.mu.RUnlock() + snapshot, waitErr := m.executionService.Wait(ctx, handle.ID, waitTimeout) + if errors.Is(waitErr, ErrExecutionWaitTimeout) { + return externalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil + } + if waitErr != nil { + return nil, handle.ID, waitErr + } + if snapshot == nil || snapshot.Execution == nil { + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回执行快照"}}, IsError: true}, handle.ID, nil + } + if snapshot.Execution.Result != nil { + return snapshot.Execution.Result, handle.ID, nil + } + if snapshot.Execution.Error != "" { + return nil, handle.ID, errors.New(snapshot.Execution.Error) + } + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}, IsError: false}, handle.ID, nil +} + +func externalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult { + execID := "" + status := ToolExecutionStatusRunning + toolName := "" + elapsed := time.Duration(0) + if snapshot != nil && snapshot.Execution != nil { + execID = snapshot.Execution.ID + status = snapshot.Execution.Status + toolName = snapshot.Execution.ToolName + elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second) + } + waitText := "unbounded" + if waitTimeout > 0 { + waitText = waitTimeout.Round(time.Second).String() + } + msg := fmt.Sprintf(`工具已提交到后台执行，但本次等待已到达上限。 + +execution_id: %s +tool: %s +status: %s +wait_timeout: %s +elapsed: %s + +你可以继续推理、改用其他工具，或调用 wait_tool_execution 继续等待该 execution_id；也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed) + return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true} +} + +func (m *ExternalMCPManager) checkExternalMCPCircuit(mcpName string) error { + if m == nil { + return nil + } + name := strings.TrimSpace(mcpName) + if name == "" { + return nil + } + m.mu.Lock() + defer m.mu.Unlock() + if m.resilience.CircuitFailureThreshold < 0 { + return nil + } + rt := m.externalMCPRuntimeLocked(name) + if rt == nil || rt.circuitOpenUntil.IsZero() { + return nil + } + now := time.Now() + if now.Before(rt.circuitOpenUntil) { + return fmt.Errorf("外部MCP服务 %s 已临时熔断，预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second)) + } + rt.circuitOpenUntil = time.Time{} + return nil +} + +func (m *ExternalMCPManager) acquireExternalMCPCallSlot(ctx context.Context, mcpName string) (func(), error) { + if m == nil { + return func() {}, nil + } + name := strings.TrimSpace(mcpName) + m.mu.Lock() + rt := m.externalMCPRuntimeLocked(name) + serverSem := chan struct{}(nil) + if rt != nil { + serverSem = rt.semaphore + } + globalSem := m.globalSemaphore + m.mu.Unlock() + + releaseGlobal := false + if globalSem != nil { + select { + case globalSem <- struct{}{}: + releaseGlobal = true + case <-ctxDone(ctx): + return func() {}, contextErr(ctx) + } + } + releaseServer := false + if serverSem != nil { + select { + case serverSem <- struct{}{}: + releaseServer = true + case <-ctxDone(ctx): + if releaseGlobal { + <-globalSem + } + return func() {}, contextErr(ctx) + } + } + return func() { + if releaseServer { + <-serverSem + } + if releaseGlobal { + <-globalSem + } + }, nil +} + +func contextErr(ctx context.Context) error { + if ctx == nil || ctx.Err() == nil { + return context.Canceled + } + return ctx.Err() +} + +func (m *ExternalMCPManager) recordExternalMCPResult(mcpName string, failed bool) { + if m == nil || strings.TrimSpace(mcpName) == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + rt := m.externalMCPRuntimeLocked(mcpName) + if rt == nil { + return + } + if !failed { + rt.consecutiveFailures = 0 + rt.circuitOpenUntil = time.Time{} + return + } + if m.resilience.CircuitFailureThreshold < 0 { + return + } + rt.consecutiveFailures++ + if rt.consecutiveFailures >= m.resilience.CircuitFailureThreshold { + rt.circuitOpenUntil = time.Now().Add(m.resilience.CircuitCooldown) + m.logger.Warn("外部MCP服务触发熔断", + zap.String("name", mcpName), + zap.Int("consecutiveFailures", rt.consecutiveFailures), + zap.Duration("cooldown", m.resilience.CircuitCooldown), + ) + } +} + +func (m *ExternalMCPManager) externalMCPRuntimeLocked(mcpName string) *externalMCPServerRuntime { + if m.serverRuntimes == nil { + m.serverRuntimes = make(map[string]*externalMCPServerRuntime) + } + name := strings.TrimSpace(mcpName) + if name == "" { + return nil + } + if rt := m.serverRuntimes[name]; rt != nil { + return rt + } + var sem chan struct{} + if m.resilience.MaxConcurrentPerServer > 0 { + sem = make(chan struct{}, m.resilience.MaxConcurrentPerServer) + } + rt := &externalMCPServerRuntime{semaphore: sem} + m.serverRuntimes[name] = rt + return rt +} + +func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(m.readAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + _ = m.takeAbortUserNote(executionID) + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + *err = nil + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + return true +} + +func (m *ExternalMCPManager) readAbortUserNote(id string) string { + m.mu.Lock() + defer m.mu.Unlock() + if m.abortUserNotes == nil { + return "" + } + return m.abortUserNotes[id] +} + +func (m *ExternalMCPManager) takeAbortUserNote(id string) string { + m.mu.Lock() + defer m.mu.Unlock() + if m.abortUserNotes == nil { + return "" + } + n := m.abortUserNotes[id] + delete(m.abortUserNotes, id) + return n +} + +// cleanupOldExecutions 清理旧的执行记录（保持内存中的记录数量在限制内） +func (m *ExternalMCPManager) cleanupOldExecutions() { + const maxExecutionsInMemory = 1000 + if len(m.executions) <= maxExecutionsInMemory { + return + } + + // 按开始时间排序，删除最旧的记录 + type execTime struct { + id string + startTime time.Time + } + var execs []execTime + for id, exec := range m.executions { + execs = append(execs, execTime{id: id, startTime: exec.StartTime}) + } + + // 按时间排序 + for i := 0; i < len(execs)-1; i++ { + for j := i + 1; j < len(execs); j++ { + if execs[i].startTime.After(execs[j].startTime) { + execs[i], execs[j] = execs[j], execs[i] + } + } + } + + // 删除最旧的记录 + toDelete := len(m.executions) - maxExecutionsInMemory + for i := 0; i < toDelete && i < len(execs); i++ { + delete(m.executions, execs[i].id) + } +} + +// GetExecution 获取执行记录（先从内存查找，再从数据库查找） +func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) { + if m.executionService != nil { + if snap, err := m.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil { + return snap.Execution, true + } + } + m.mu.RLock() + exec, exists := m.executions[id] + m.mu.RUnlock() + + if exists { + return exec, true + } + + if m.storage != nil { + exec, err := m.storage.GetToolExecution(id) + if err == nil { + return exec, true + } + } + + return nil, false +} + +func (m *ExternalMCPManager) registerRunningCancel(id string, cancel context.CancelFunc) { + m.mu.Lock() + m.runningCancels[id] = cancel + m.mu.Unlock() +} + +func (m *ExternalMCPManager) unregisterRunningCancel(id string) { + m.mu.Lock() + delete(m.runningCancels, id) + m.mu.Unlock() +} + +// CancelToolExecutionWithNote 取消外部 MCP 工具；note 非空时与已返回输出合并后交给模型。 +func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool { + if m.executionService != nil && m.executionService.Cancel(id, note) { + return true + } + m.mu.Lock() + cancel, ok := m.runningCancels[id] + if !ok || cancel == nil { + m.mu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + if m.abortUserNotes == nil { + m.abortUserNotes = make(map[string]string) + } + m.abortUserNotes[id] = strings.TrimSpace(note) + } + m.mu.Unlock() + cancel() + return true +} + +// CancelToolExecution 取消正在执行的外部 MCP 工具（无用户说明）。 +func (m *ExternalMCPManager) CancelToolExecution(id string) bool { + return m.CancelToolExecutionWithNote(id, "") +} + +// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的外部 MCP executionId 快照。 +func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} { + if m == nil { + return nil + } + if m.executionService != nil { + if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 { + return ids + } + } + m.mu.Lock() + defer m.mu.Unlock() + if len(m.runningCancels) == 0 { + return nil + } + out := make(map[string]struct{}, len(m.runningCancels)) + for id := range m.runningCancels { + out[id] = struct{}{} + } + return out +} + +// updateStats 更新统计信息 +func (m *ExternalMCPManager) updateStats(toolName string, failed bool) { + now := time.Now() + if m.storage != nil { + totalCalls := 1 + successCalls := 0 + failedCalls := 0 + if failed { + failedCalls = 1 + } else { + successCalls = 1 + } + if err := m.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil { + m.logger.Warn("保存统计信息到数据库失败", zap.Error(err)) + } + return + } + + m.mu.Lock() + defer m.mu.Unlock() + + if m.stats[toolName] == nil { + m.stats[toolName] = &ToolStats{ + ToolName: toolName, + } + } + + stats := m.stats[toolName] + stats.TotalCalls++ + stats.LastCallTime = &now + + if failed { + stats.FailedCalls++ + } else { + stats.SuccessCalls++ + } +} + +// GetStats 获取MCP服务器统计信息 +func (m *ExternalMCPManager) GetStats() map[string]interface{} { + m.mu.RLock() + defer m.mu.RUnlock() + + total := len(m.configs) + enabled := 0 + disabled := 0 + connected := 0 + + for name, cfg := range m.configs { + if m.isEnabled(cfg) { + enabled++ + if client, exists := m.clients[name]; exists && client.IsConnected() { + connected++ + } + } else { + disabled++ + } + } + + return map[string]interface{}{ + "total": total, + "enabled": enabled, + "disabled": disabled, + "connected": connected, + } +} + +// GetToolStats 获取工具统计信息（合并内存和数据库） +// 只返回外部MCP工具的统计信息（工具名称包含 "::"） +func (m *ExternalMCPManager) GetToolStats() map[string]*ToolStats { + result := make(map[string]*ToolStats) + + // 从数据库加载统计信息（如果使用数据库存储） + if m.storage != nil { + dbStats, err := m.storage.LoadToolStats() + if err == nil { + // 只保留外部MCP工具的统计信息（工具名称包含 "::"） + for k, v := range dbStats { + if findSubstring(k, "::") > 0 { + result[k] = v + } + } + } else { + m.logger.Warn("从数据库加载统计信息失败", zap.Error(err)) + } + } + + // 合并内存中的统计信息 + m.mu.RLock() + for k, v := range m.stats { + // 如果数据库中已有该工具的统计信息，合并它们 + if existing, exists := result[k]; exists { + // 创建新的统计信息对象，避免修改共享对象 + merged := &ToolStats{ + ToolName: k, + TotalCalls: existing.TotalCalls + v.TotalCalls, + SuccessCalls: existing.SuccessCalls + v.SuccessCalls, + FailedCalls: existing.FailedCalls + v.FailedCalls, + } + // 使用最新的调用时间 + if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) { + merged.LastCallTime = v.LastCallTime + } else if existing.LastCallTime != nil { + timeCopy := *existing.LastCallTime + merged.LastCallTime = &timeCopy + } + result[k] = merged + } else { + // 如果数据库中没有，直接使用内存中的统计信息 + statCopy := *v + result[k] = &statCopy + } + } + m.mu.RUnlock() + + return result +} + +// GetToolCount 获取指定外部MCP的工具数量（从缓存读取，不阻塞） +func (m *ExternalMCPManager) GetToolCount(name string) (int, error) { + // 先从缓存读取 + m.toolCountsMu.RLock() + if count, exists := m.toolCounts[name]; exists { + m.toolCountsMu.RUnlock() + return count, nil + } + m.toolCountsMu.RUnlock() + + // 如果缓存中没有，检查客户端状态 + client, exists := m.GetClient(name) + if !exists { + return 0, fmt.Errorf("客户端不存在: %s", name) + } + + if !client.IsConnected() { + // 未连接，缓存为0 + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() + return 0, nil + } + + // 如果已连接但缓存中没有，触发异步刷新并返回0（避免阻塞） + m.triggerToolCountRefresh() + return 0, nil +} + +// GetToolCounts 获取所有外部MCP的工具数量（从缓存读取，不阻塞） +func (m *ExternalMCPManager) GetToolCounts() map[string]int { + m.toolCountsMu.RLock() + defer m.toolCountsMu.RUnlock() + + // 返回缓存的副本，避免外部修改 + result := make(map[string]int) + for k, v := range m.toolCounts { + result[k] = v + } + return result +} + +// refreshToolCounts 刷新工具数量缓存（后台异步执行） +// 使用 atomic flag 防止并发堆积：如果上一次刷新尚未完成，本次触发直接跳过。 +func (m *ExternalMCPManager) refreshToolCounts() { + if !m.refreshing.CompareAndSwap(false, true) { + return // 上一次刷新尚未完成，跳过 + } + defer m.refreshing.Store(false) + + m.mu.RLock() + clients := make(map[string]ExternalMCPClient) + for k, v := range m.clients { + clients[k] = v + } + m.mu.RUnlock() + + newCounts := make(map[string]int) + + // 使用goroutine并发获取每个客户端的工具数量，避免串行阻塞 + type countResult struct { + name string + count int + } + resultChan := make(chan countResult, len(clients)) + + for name, client := range clients { + go func(n string, c ExternalMCPClient) { + if !c.IsConnected() { + resultChan <- countResult{name: n, count: 0} + return + } + + // 缓存仍新鲜时直接复用，避免与 GetAllTools 重复打远程 + if _, fresh := m.getFreshCachedTools(n); fresh { + m.toolCountsMu.RLock() + count := m.toolCounts[n] + m.toolCountsMu.RUnlock() + resultChan <- countResult{name: n, count: count} + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + tools, err := m.listToolsDeduped(ctx, n, c) + cancel() + + if err != nil { + if !isConnectionDeadError(err) { + m.logger.Warn("获取外部MCP工具数量失败，请检查连接或服务端 tools/list", + zap.String("name", n), + zap.Error(err), + ) + } + resultChan <- countResult{name: n, count: -1} + return + } + + resultChan <- countResult{name: n, count: len(tools)} + }(name, client) + } + + // 收集结果 + m.toolCountsMu.RLock() + oldCounts := make(map[string]int) + for k, v := range m.toolCounts { + oldCounts[k] = v + } + m.toolCountsMu.RUnlock() + + for i := 0; i < len(clients); i++ { + result := <-resultChan + if result.count >= 0 { + newCounts[result.name] = result.count + } else { + // 获取失败，保留旧值 + if oldCount, exists := oldCounts[result.name]; exists { + newCounts[result.name] = oldCount + } else { + newCounts[result.name] = 0 + } + } + } + + // 更新缓存 + m.toolCountsMu.Lock() + // 更新所有获取到的值 + for name, count := range newCounts { + m.toolCounts[name] = count + } + // 对于未连接的客户端，设置为0 + for name, client := range clients { + if !client.IsConnected() { + m.toolCounts[name] = 0 + } + } + m.toolCountsMu.Unlock() +} + +// refreshToolCache 刷新指定MCP的工具列表缓存 +func (m *ExternalMCPManager) refreshToolCache(name string, client ExternalMCPClient) { + if !client.IsConnected() { + return + } + if client.GetStatus() == "error" { + m.logger.Debug("跳过刷新工具列表缓存（连接失败）", + zap.String("name", name), + ) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if _, err := m.listToolsDeduped(ctx, name, client); err != nil { + m.logger.Debug("刷新工具列表缓存失败", + zap.String("name", name), + zap.Error(err), + ) + } +} + +// startToolCountRefresh 启动后台刷新工具数量的goroutine +func (m *ExternalMCPManager) startToolCountRefresh() { + m.refreshWg.Add(1) + go func() { + defer m.refreshWg.Done() + ticker := time.NewTicker(externalToolCountRefreshInterval) + defer ticker.Stop() + + // 立即执行一次刷新 + m.refreshToolCounts() + + for { + select { + case <-ticker.C: + m.refreshToolCounts() + case <-m.stopRefresh: + return + } + } + }() +} + +// triggerToolCountRefresh 触发立即刷新工具数量（异步） +func (m *ExternalMCPManager) triggerToolCountRefresh() { + go m.refreshToolCounts() +} + +// createClient 创建客户端（不连接）。统一使用官方 MCP Go SDK 的 lazy 客户端，连接在 Initialize 时完成。 +func (m *ExternalMCPManager) createClient(serverCfg config.ExternalMCPServerConfig) ExternalMCPClient { + transport := serverCfg.GetTransportType() + + switch transport { + case "http": + if serverCfg.URL == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + case "stdio": + if serverCfg.Command == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + case "sse": + if serverCfg.URL == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + default: + if transport == "" { + return nil + } + // 未知传输类型也尝试使用 lazy client + return newLazySDKClient(serverCfg, m.logger) + } +} + +// doConnect 执行实际连接 +func (m *ExternalMCPManager) doConnect(name string, serverCfg config.ExternalMCPServerConfig, client ExternalMCPClient) error { + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + // 初始化连接 + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + if err := client.Initialize(ctx); err != nil { + return err + } + + m.logger.Info("外部MCP客户端已连接", + zap.String("name", name), + ) + + return nil +} + +// setClientStatus 设置客户端状态（通过类型断言） +func (m *ExternalMCPManager) setClientStatus(client ExternalMCPClient, status string) { + if c, ok := client.(*lazySDKClient); ok { + c.setStatus(status) + } +} + +// connectClient 连接客户端（异步）- 保留用于向后兼容 +func (m *ExternalMCPManager) connectClient(name string, serverCfg config.ExternalMCPServerConfig) error { + client := m.createClient(serverCfg) + if client == nil { + return fmt.Errorf("无法创建客户端：不支持的传输模式") + } + + // 设置状态为connecting + m.setClientStatus(client, "connecting") + + // 初始化连接 + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + if err := client.Initialize(ctx); err != nil { + m.logger.Error("初始化外部MCP客户端失败", + zap.String("name", name), + zap.Error(err), + ) + return err + } + + // 保存客户端 + m.mu.Lock() + m.clients[name] = client + m.mu.Unlock() + + m.logger.Info("外部MCP客户端已连接", + zap.String("name", name), + ) + + m.onClientConnected(name) + + // 连接成功，触发工具数量刷新和工具列表缓存刷新 + m.triggerToolCountRefresh() + m.mu.RLock() + if client, exists := m.clients[name]; exists { + m.refreshToolCache(name, client) + } + m.mu.RUnlock() + + return nil +} + +// isEnabled 检查是否启用 +func (m *ExternalMCPManager) isEnabled(cfg config.ExternalMCPServerConfig) bool { + return cfg.ExternalMCPEnable +} + +// findSubstring 查找子字符串（简单实现） +func findSubstring(s, substr string) int { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return i + } + } + return -1 +} + +// StartAllEnabled 启动所有启用的客户端 +func (m *ExternalMCPManager) StartAllEnabled() { + m.mu.RLock() + configs := make(map[string]config.ExternalMCPServerConfig) + for k, v := range m.configs { + configs[k] = v + } + m.mu.RUnlock() + + for name, cfg := range configs { + if m.isEnabled(cfg) { + go func(n string, c config.ExternalMCPServerConfig) { + if err := m.connectClient(n, c); err != nil { + // 检查是否是连接被拒绝的错误（服务可能还没启动） + errStr := strings.ToLower(err.Error()) + isConnectionRefused := strings.Contains(errStr, "connection refused") || + strings.Contains(errStr, "dial tcp") || + strings.Contains(errStr, "connect: connection refused") + + if isConnectionRefused { + // 连接被拒绝，说明目标服务可能还没启动，这是正常的 + // 使用 Warn 级别，提示用户这是正常的，可以通过手动启动或等待服务启动后自动连接 + fields := []zap.Field{ + zap.String("name", n), + zap.String("message", "目标服务可能尚未启动，这是正常的。服务启动后可通过界面手动连接，或等待自动重试"), + zap.Error(err), + } + + transport := c.GetTransportType() + + if transport == "http" && c.URL != "" { + fields = append(fields, zap.String("url", c.URL)) + } else if transport == "stdio" && c.Command != "" { + fields = append(fields, zap.String("command", c.Command)) + } + + m.logger.Warn("外部MCP服务暂未就绪", fields...) + } else { + // 其他错误，使用 Error 级别 + m.logger.Error("启动外部MCP客户端失败", + zap.String("name", n), + zap.Error(err), + ) + } + } + }(name, cfg) + } + } +} + +// StopAll 停止所有客户端 +func (m *ExternalMCPManager) StopAll() { + if m.executionService != nil { + m.executionService.CancelAll("外部 MCP 管理器正在停止") + } + clients := make(map[string]ExternalMCPClient) + m.mu.Lock() + for name, client := range m.clients { + clients[name] = client + delete(m.clients, name) + } + m.mu.Unlock() + + for name, client := range clients { + if client != nil { + _ = client.Close() + } + m.clearReconnectState(name) + } + + // 清理所有工具数量缓存 + m.toolCountsMu.Lock() + m.toolCounts = make(map[string]int) + m.toolCountsMu.Unlock() + + // 清理所有工具列表缓存 + m.toolCacheMu.Lock() + m.toolCache = make(map[string]toolListCacheEntry) + m.toolCacheMu.Unlock() + + // 停止后台刷新（使用 select 避免重复关闭 channel） + select { + case <-m.stopRefresh: + // 已经关闭，不需要再次关闭 + default: + close(m.stopRefresh) + } + m.refreshWg.Wait() +} diff --git a/internal/mcp/external_manager_async_test.go b/internal/mcp/external_manager_async_test.go new file mode 100644 index 00000000..6c280ed9 --- /dev/null +++ b/internal/mcp/external_manager_async_test.go @@ -0,0 +1,230 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "go.uber.org/zap" +) + +type blockingExternalMCPClient struct { + started chan struct{} + calls chan string + release chan struct{} + result *ToolResult + count atomic.Int32 +} + +func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient { + return &blockingExternalMCPClient{ + started: make(chan struct{}), + calls: make(chan string, 8), + release: make(chan struct{}), + result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}, + } +} + +func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil } +func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) { + return []Tool{{Name: "slow_tool"}}, nil +} +func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + c.count.Add(1) + select { + case c.calls <- name: + default: + } + select { + case <-c.started: + default: + close(c.started) + } + select { + case <-c.release: + return c.result, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} +func (c *blockingExternalMCPClient) Close() error { return nil } +func (c *blockingExternalMCPClient) IsConnected() bool { return true } +func (c *blockingExternalMCPClient) GetStatus() string { return "connected" } + +type failingExternalMCPClient struct{} + +func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil } +func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) { + return []Tool{{Name: "fail_tool"}}, nil +} +func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + return nil, errors.New("boom") +} +func (c *failingExternalMCPClient) Close() error { return nil } +func (c *failingExternalMCPClient) IsConnected() bool { return true } +func (c *failingExternalMCPClient) GetStatus() string { return "connected" } + +func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.ConfigureToolWaitTimeoutSeconds(1) + manager.toolWaitTimeout = 10 * time.Millisecond + client := newBlockingExternalMCPClient("slow result ready") + manager.clients["lab"] = client + + callCtx, callCancel := context.WithCancel(context.Background()) + result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"}) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if executionID == "" { + t.Fatal("expected execution id") + } + if result == nil || !result.IsError { + t.Fatalf("expected soft timeout tool result, got %#v", result) + } + text := ToolResultPlainText(result) + if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") { + t.Fatalf("timeout result should include execution id and wait guidance, got %q", text) + } + + select { + case <-client.started: + default: + t.Fatal("worker did not start") + } + callCancel() + close(client.release) + + snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second) + if err != nil { + t.Fatalf("Wait returned error: %v", err) + } + if snapshot == nil || snapshot.Execution == nil { + t.Fatal("expected execution snapshot") + } + if snapshot.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("status = %q, want completed", snapshot.Execution.Status) + } + if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" { + t.Fatalf("result = %q, want slow result ready", got) + } +} + +func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.toolWaitTimeout = 10 * time.Millisecond + client := newBlockingExternalMCPClient("control wait result") + manager.clients["lab"] = client + + result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID) + } + + server := NewServer(zap.NewNop()) + RegisterExecutionControlTools(server, manager) + close(client.release) + + waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 1, + }) + if err != nil { + t.Fatalf("wait_tool_execution returned error: %v", err) + } + if waitResult == nil || waitResult.IsError { + t.Fatalf("expected successful wait result, got %#v", waitResult) + } + body := ToolResultPlainText(waitResult) + if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") { + t.Fatalf("wait result body missing completed status/result: %s", body) + } +} + +func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.toolWaitTimeout = 10 * time.Millisecond + manager.ConfigureResilience(ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 1, + MaxConcurrentTotal: 4, + CircuitFailureThreshold: -1, + CircuitCooldown: time.Second, + }) + client := newBlockingExternalMCPClient("ok") + manager.clients["lab"] = client + + done1 := make(chan struct{}) + go func() { + _, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil) + close(done1) + }() + select { + case <-client.calls: + case <-time.After(time.Second): + t.Fatal("first worker did not enter client") + } + + type callOutcome struct { + executionID string + err error + } + done2 := make(chan callOutcome, 1) + go func() { + _, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil) + done2 <- callOutcome{executionID: executionID, err: err} + }() + select { + case <-client.calls: + t.Fatal("second worker entered client before per-server slot was released") + case <-time.After(50 * time.Millisecond): + } + var second callOutcome + select { + case second = <-done2: + case <-time.After(time.Second): + t.Fatal("second call did not return after bounded wait") + } + if second.err != nil || second.executionID == "" { + t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err) + } + snapshot, err := manager.executionService.Get(second.executionID) + if err != nil { + t.Fatalf("Get queued execution: %v", err) + } + if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued { + t.Fatalf("second execution status = %#v, want queued", snapshot) + } + close(client.release) + select { + case <-client.calls: + case <-time.After(time.Second): + t.Fatal("second worker did not enter client after slot release") + } + <-done1 +} + +func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.ConfigureResilience(ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 2, + MaxConcurrentTotal: 4, + CircuitFailureThreshold: 1, + CircuitCooldown: time.Minute, + }) + manager.clients["lab"] = &failingExternalMCPClient{} + + _, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("expected first call to fail with client error, got %v", err) + } + _, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil) + if err == nil || !strings.Contains(err.Error(), "熔断") { + t.Fatalf("expected circuit breaker rejection, got %v", err) + } +} diff --git a/internal/mcp/external_manager_test.go b/internal/mcp/external_manager_test.go new file mode 100644 index 00000000..3baff567 --- /dev/null +++ b/internal/mcp/external_manager_test.go @@ -0,0 +1,261 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + t.Cleanup(manager.StopAll) + manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error { + return errors.New("denied by policy") + }) + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true})) + _, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{}) + if err == nil || !strings.Contains(err.Error(), "authorization denied") { + t.Fatalf("external call bypassed authorizer: %v", err) + } + if executionID == "" { + t.Fatal("denied external call should still return an execution id") + } + execution, ok := manager.GetExecution(executionID) + if !ok || execution == nil { + t.Fatalf("missing denied external execution %q", executionID) + } + if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") { + t.Fatalf("denied external execution = %#v, want failed with policy error", execution) + } +} + +func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 测试添加stdio配置 + stdioCfg := config.ExternalMCPServerConfig{ + Command: "python3", + Args: []string{"/path/to/script.py"}, + Description: "Test stdio MCP", + Timeout: 30, + ExternalMCPEnable: true, + } + + err := manager.AddOrUpdateConfig("test-stdio", stdioCfg) + if err != nil { + t.Fatalf("添加stdio配置失败: %v", err) + } + + // 测试添加HTTP配置 + httpCfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://127.0.0.1:8081/mcp", + Description: "Test HTTP MCP", + Timeout: 30, + ExternalMCPEnable: false, + } + + err = manager.AddOrUpdateConfig("test-http", httpCfg) + if err != nil { + t.Fatalf("添加HTTP配置失败: %v", err) + } + + // 验证配置已保存 + configs := manager.GetConfigs() + if len(configs) != 2 { + t.Fatalf("期望2个配置，实际%d个", len(configs)) + } + + if configs["test-stdio"].Command != stdioCfg.Command { + t.Errorf("stdio配置命令不匹配") + } + + if configs["test-http"].URL != httpCfg.URL { + t.Errorf("HTTP配置URL不匹配") + } +} + +func TestExternalMCPManager_RemoveConfig(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + cfg := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + } + + manager.AddOrUpdateConfig("test-remove", cfg) + + // 移除配置 + err := manager.RemoveConfig("test-remove") + if err != nil { + t.Fatalf("移除配置失败: %v", err) + } + + configs := manager.GetConfigs() + if _, exists := configs["test-remove"]; exists { + t.Error("配置应该已被移除") + } +} + +func TestExternalMCPManager_GetStats(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 添加多个配置 + manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + }) + + manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{ + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: true, + }) + + manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + }) + + stats := manager.GetStats() + + if stats["total"].(int) != 3 { + t.Errorf("期望总数3，实际%d", stats["total"]) + } + + if stats["enabled"].(int) != 2 { + t.Errorf("期望启用数2，实际%d", stats["enabled"]) + } + + if stats["disabled"].(int) != 1 { + t.Errorf("期望停用数1，实际%d", stats["disabled"]) + } +} + +func TestExternalMCPManager_LoadConfigs(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + externalMCPConfig := config.ExternalMCPConfig{ + Servers: map[string]config.ExternalMCPServerConfig{ + "loaded1": { + Command: "python3", + ExternalMCPEnable: true, + }, + "loaded2": { + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: false, + }, + }, + } + + manager.LoadConfigs(&externalMCPConfig) + + configs := manager.GetConfigs() + if len(configs) != 2 { + t.Fatalf("期望2个配置，实际%d个", len(configs)) + } + + if configs["loaded1"].Command != "python3" { + t.Error("配置1加载失败") + } + + if configs["loaded2"].URL != "http://127.0.0.1:8081/mcp" { + t.Error("配置2加载失败") + } +} + +// TestLazySDKClient_InitializeFails 验证无效配置时 SDK 客户端 Initialize 失败并设置 error 状态 +func TestLazySDKClient_InitializeFails(t *testing.T) { + logger := zap.NewNop() + // 使用不存在的 HTTP 地址，Initialize 应失败 + cfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://127.0.0.1:19999/nonexistent", + Timeout: 2, + } + c := newLazySDKClient(cfg, logger) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := c.Initialize(ctx) + if err == nil { + t.Fatal("expected error when connecting to invalid server") + } + if c.GetStatus() != "error" { + t.Errorf("expected status error, got %s", c.GetStatus()) + } + c.Close() +} + +func TestExternalMCPManager_StartStopClient(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 添加一个禁用的配置 + cfg := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + } + + manager.AddOrUpdateConfig("test-start-stop", cfg) + + // 尝试启动（可能会失败，因为没有真实的服务器） + err := manager.StartClient("test-start-stop") + if err != nil { + t.Logf("启动失败（可能是没有服务器）: %v", err) + } + + // 停止 + err = manager.StopClient("test-start-stop") + if err != nil { + t.Fatalf("停止失败: %v", err) + } + + // 验证配置已更新为禁用 + configs := manager.GetConfigs() + if configs["test-start-stop"].ExternalMCPEnable { + t.Error("配置应该已被禁用") + } +} + +func TestExternalMCPManager_CallTool(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 测试调用不存在的工具 + _, _, err := manager.CallTool(context.Background(), "nonexistent::tool", map[string]interface{}{}) + if err == nil { + t.Error("应该返回错误") + } + + // 测试无效的工具名称格式 + _, _, err = manager.CallTool(context.Background(), "invalid-tool-name", map[string]interface{}{}) + if err == nil { + t.Error("应该返回错误（无效格式）") + } +} + +func TestExternalMCPManager_GetAllTools(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + ctx := context.Background() + tools, err := manager.GetAllTools(ctx) + if err != nil { + t.Fatalf("获取工具列表失败: %v", err) + } + + // 如果没有连接的客户端，应该返回空列表 + if len(tools) != 0 { + t.Logf("获取到%d个工具", len(tools)) + } +} diff --git a/internal/mcp/run_context.go b/internal/mcp/run_context.go new file mode 100644 index 00000000..7612032e --- /dev/null +++ b/internal/mcp/run_context.go @@ -0,0 +1,147 @@ +package mcp + +import ( + "context" + "strings" +) + +// ToolRunRegistry 在工具开始/结束时登记当前 executionId，供对话页「仅终止当前工具」与监控页共用取消逻辑。 +type ToolRunRegistry interface { + RegisterRunningTool(conversationID, executionID string) + UnregisterRunningTool(conversationID, executionID string) +} + +// EinoExecuteRunRegistry 登记进行中的 Eino filesystem execute，供「中断并继续」终止 amass 等长命令。 +type EinoExecuteRunRegistry interface { + RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc) + UnregisterActiveEinoExecute(conversationID string) + AbortActiveEinoExecute(conversationID, note string) bool + TakeEinoExecuteAbortNote(conversationID string) string +} + +type toolRunRegistryCtxKey struct{} +type einoExecuteRunRegistryCtxKey struct{} +type mcpConversationIDCtxKey struct{} +type mcpExecutionIDCtxKey struct{} +type mcpProjectIDCtxKey struct{} + +// WithToolRunRegistry 将登记器注入 ctx（Eino / 原生 Agent 任务 ctx）。 +func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context { + if ctx == nil || reg == nil { + return ctx + } + return context.WithValue(ctx, toolRunRegistryCtxKey{}, reg) +} + +// ToolRunRegistryFromContext 取出登记器（无则 nil）。 +func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry { + if ctx == nil { + return nil + } + v, _ := ctx.Value(toolRunRegistryCtxKey{}).(ToolRunRegistry) + return v +} + +// WithEinoExecuteRunRegistry 将 Eino execute 取消登记器注入 ctx。 +func WithEinoExecuteRunRegistry(ctx context.Context, reg EinoExecuteRunRegistry) context.Context { + if ctx == nil || reg == nil { + return ctx + } + return context.WithValue(ctx, einoExecuteRunRegistryCtxKey{}, reg) +} + +// EinoExecuteRunRegistryFromContext 取出 Eino execute 登记器（无则 nil）。 +func EinoExecuteRunRegistryFromContext(ctx context.Context) EinoExecuteRunRegistry { + if ctx == nil { + return nil + } + v, _ := ctx.Value(einoExecuteRunRegistryCtxKey{}).(EinoExecuteRunRegistry) + return v +} + +// WithMCPConversationID 将对话 ID 注入 ctx，供 CallTool 内与 executionId 关联。 +func WithMCPConversationID(ctx context.Context, conversationID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(conversationID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpConversationIDCtxKey{}, id) +} + +// MCPConversationIDFromContext 读取对话 ID。 +func MCPConversationIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpConversationIDCtxKey{}).(string) + return v +} + +// WithMCPExecutionID 将当前工具 executionId 注入 ctx，供超长输出落盘文件名对齐。 +func WithMCPExecutionID(ctx context.Context, executionID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(executionID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id) +} + +// MCPExecutionIDFromContext 读取当前工具 executionId。 +func MCPExecutionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string) + return v +} + +// WithMCPProjectID 将项目 ID 注入 ctx，供 reduction/trunc 落盘路径与项目隔离对齐。 +func WithMCPProjectID(ctx context.Context, projectID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(projectID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpProjectIDCtxKey{}, id) +} + +// MCPProjectIDFromContext 读取项目 ID。 +func MCPProjectIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string) + return v +} + +func notifyToolRunBegin(ctx context.Context, executionID string) { + reg := ToolRunRegistryFromContext(ctx) + if reg == nil { + return + } + conv := MCPConversationIDFromContext(ctx) + if conv == "" || strings.TrimSpace(executionID) == "" { + return + } + reg.RegisterRunningTool(conv, executionID) +} + +func notifyToolRunEnd(ctx context.Context, executionID string) { + reg := ToolRunRegistryFromContext(ctx) + if reg == nil { + return + } + conv := MCPConversationIDFromContext(ctx) + if conv == "" || strings.TrimSpace(executionID) == "" { + return + } + reg.UnregisterRunningTool(conv, executionID) +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go new file mode 100644 index 00000000..3760183b --- /dev/null +++ b/internal/mcp/server.go @@ -0,0 +1,1704 @@ +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "sort" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/mcp/builtin" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// MonitorStorage 监控数据存储接口 +type MonitorStorage interface { + SaveToolExecution(exec *ToolExecution) error + UpdateToolExecutionResult(id string, result *ToolResult) error + LoadToolExecutions() ([]*ToolExecution, error) + GetToolExecution(id string) (*ToolExecution, error) + SaveToolStats(toolName string, stats *ToolStats) error + LoadToolStats() (map[string]*ToolStats, error) + UpdateToolStats(toolName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error +} + +// Server MCP服务器 +type Server struct { + tools map[string]ToolHandler + toolDefs map[string]Tool // 工具定义 + executions map[string]*ToolExecution + stats map[string]*ToolStats + prompts map[string]*Prompt // 提示词模板 + resources map[string]*Resource // 资源 + storage MonitorStorage // 可选的持久化存储 + mu sync.RWMutex + logger *zap.Logger + maxExecutionsInMemory int // 内存中最大执行记录数 + sseClients map[string]*sseClient + runningCancels map[string]context.CancelFunc + runningCancelsMu sync.Mutex + abortUserNotes map[string]string // 监控页终止时附带的用户说明，与 executionID 对应 + // httpToolTimeoutMinutes 同步 agent.tool_timeout_minutes，用于 POST /api/mcp 的 tools/call（不经 Agent 包装的路径）。 + // nil 表示未配置，沿用默认 30 分钟；指向 0 表示不限制；>0 为分钟数。 + httpToolTimeoutMinutes *int + httpToolTimeoutMu sync.RWMutex + toolAuthorizer func(context.Context, string, map[string]interface{}) error + executionService *ExecutionService + toolWaitTimeout time.Duration + toolResultMaxBytes int + spillRootDir string +} + +const defaultPartialOutputMaxBytes = 64 * 1024 + +// SetToolAuthorizer installs the common policy decision point for every +// user-attributed tool call, whether it originates from HTTP or an Agent. +func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + if s == nil { + return + } + s.mu.Lock() + s.toolAuthorizer = authorizer + s.mu.Unlock() +} + +type sseClient struct { + id string + send chan []byte +} + +// ToolHandler 工具处理函数 +type ToolHandler func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) + +func executionStatusAndMessage(err error) (status string, errMsg string) { + if errors.Is(err, context.Canceled) { + return "cancelled", "已手动终止（MCP 监控）" + } + return "failed", err.Error() +} + +// NewServer 创建新的MCP服务器 +func NewServer(logger *zap.Logger) *Server { + return NewServerWithStorage(logger, nil) +} + +// NewServerWithStorage 创建新的MCP服务器（带持久化存储） +func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server { + s := &Server{ + tools: make(map[string]ToolHandler), + toolDefs: make(map[string]Tool), + executions: make(map[string]*ToolExecution), + stats: make(map[string]*ToolStats), + prompts: make(map[string]*Prompt), + resources: make(map[string]*Resource), + storage: storage, + logger: logger, + maxExecutionsInMemory: 1000, // 默认最多在内存中保留1000条执行记录 + sseClients: make(map[string]*sseClient), + runningCancels: make(map[string]context.CancelFunc), + abortUserNotes: make(map[string]string), + toolWaitTimeout: 60 * time.Second, + toolResultMaxBytes: DefaultToolResultMaxBytes, + } + s.executionService = NewExecutionService(storage, logger) + + // 初始化默认提示词和资源 + s.initDefaultPrompts() + s.initDefaultResources() + + return s +} + +func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) { + if s == nil { + return + } + s.mu.Lock() + s.toolResultMaxBytes = maxBytes + s.mu.Unlock() + if s.executionService != nil { + s.executionService.ConfigureToolResultMaxBytes(maxBytes) + } +} + +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (s *Server) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + s.spillRootDir = strings.TrimSpace(rootDir) + s.mu.Unlock() + if s.executionService != nil { + s.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + +// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。 +// minutes<=0 表示不设置硬性截止时间（与配置「0 不限制」一致）；minutes>0 为该次调用的最长等待时间。 +// 未调用前对 tools/call 使用默认 30 分钟（与历史硬编码一致）。 +func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) { + if s == nil { + return + } + v := minutes + if v < 0 { + v = 0 + } + s.httpToolTimeoutMu.Lock() + defer s.httpToolTimeoutMu.Unlock() + s.httpToolTimeoutMinutes = &v +} + +func (s *Server) ConfigureToolWaitTimeoutSeconds(seconds int) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if seconds <= 0 { + s.toolWaitTimeout = 0 + return + } + s.toolWaitTimeout = time.Duration(seconds) * time.Second +} + +func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) { + const defaultDur = 30 * time.Minute + if parent == nil { + parent = context.Background() + } + if s == nil { + return context.WithTimeout(parent, defaultDur) + } + s.httpToolTimeoutMu.RLock() + mPtr := s.httpToolTimeoutMinutes + s.httpToolTimeoutMu.RUnlock() + if mPtr == nil { + return context.WithTimeout(parent, defaultDur) + } + if *mPtr <= 0 { + return context.WithCancel(parent) + } + return context.WithTimeout(parent, time.Duration(*mPtr)*time.Minute) +} + +// RegisterTool 注册工具 +func (s *Server) RegisterTool(tool Tool, handler ToolHandler) { + s.mu.Lock() + defer s.mu.Unlock() + s.tools[tool.Name] = handler + s.toolDefs[tool.Name] = tool + + // 自动为工具创建资源文档 + resourceURI := fmt.Sprintf("tool://%s", tool.Name) + s.resources[resourceURI] = &Resource{ + URI: resourceURI, + Name: fmt.Sprintf("%s工具文档", tool.Name), + Description: tool.Description, + MimeType: "text/plain", + } +} + +// ClearTools 清空所有工具（用于重新加载配置） +func (s *Server) ClearTools() { + s.mu.Lock() + defer s.mu.Unlock() + + // 清空工具和工具定义 + s.tools = make(map[string]ToolHandler) + s.toolDefs = make(map[string]Tool) + + // 清空工具相关的资源（保留其他资源） + newResources := make(map[string]*Resource) + for uri, resource := range s.resources { + // 保留非工具资源 + if !strings.HasPrefix(uri, "tool://") { + newResources[uri] = resource + } + } + s.resources = newResources +} + +// HandleHTTP 处理HTTP请求 +func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && strings.Contains(r.Header.Get("Accept"), "text/event-stream") { + s.handleSSE(w, r) + return + } + + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // 官方 MCP SSE 规范：带 sessionid 的 POST 表示消息发往该 SSE 会话，响应通过 SSE 流返回 + if sessionID := r.URL.Query().Get("sessionid"); sessionID != "" { + s.serveSSESessionMessage(w, r, sessionID) + return + } + + // 简单 POST：请求体为 JSON-RPC，响应在 body 中返回 + body, err := io.ReadAll(r.Body) + if err != nil { + s.sendError(w, nil, -32700, "Parse error", err.Error()) + return + } + + var msg Message + if err := json.Unmarshal(body, &msg); err != nil { + s.sendError(w, nil, -32700, "Parse error", err.Error()) + return + } + + response := s.handleMessage(r.Context(), &msg) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(response) +} + +// serveSSESessionMessage 处理发往 SSE 会话的 POST：读取 JSON-RPC 请求，处理后将响应通过该会话的 SSE 流推送 +func (s *Server) serveSSESessionMessage(w http.ResponseWriter, r *http.Request, sessionID string) { + s.mu.RLock() + client, exists := s.sseClients[sessionID] + s.mu.RUnlock() + if !exists || client == nil { + http.Error(w, "session not found", http.StatusNotFound) + return + } + + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var msg Message + if err := json.Unmarshal(body, &msg); err != nil { + http.Error(w, "failed to parse body", http.StatusBadRequest) + return + } + + response := s.handleMessage(r.Context(), &msg) + if response == nil { + w.WriteHeader(http.StatusAccepted) + return + } + + respBytes, err := json.Marshal(response) + if err != nil { + http.Error(w, "failed to encode response", http.StatusInternalServerError) + return + } + + select { + case client.send <- respBytes: + w.WriteHeader(http.StatusAccepted) + default: + http.Error(w, "session send buffer full", http.StatusServiceUnavailable) + } +} + +// handleSSE 处理 SSE 连接，兼容官方 MCP 2024-11-05 SSE 规范： +// 1. 首个事件必须为 event: endpoint，data 为客户端 POST 消息的 URL（含 sessionid） +// 2. 后续事件为 event: message，data 为 JSON-RPC 响应 +func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) { + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "Streaming unsupported", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + + sessionID := uuid.New().String() + client := &sseClient{ + id: sessionID, + send: make(chan []byte, 32), + } + + s.addSSEClient(client) + defer s.removeSSEClient(client.id) + + // 官方规范：首个事件为 endpoint，data 为消息端点 URL（客户端将向该 URL POST 请求） + scheme := "http" + if r.TLS != nil { + scheme = "https" + } + if r.URL.Scheme != "" { + scheme = r.URL.Scheme + } + endpointURL := fmt.Sprintf("%s://%s%s?sessionid=%s", scheme, r.Host, r.URL.Path, sessionID) + fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpointURL) + flusher.Flush() + + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + + for { + select { + case <-r.Context().Done(): + return + case msg, ok := <-client.send: + if !ok { + return + } + fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg) + flusher.Flush() + case <-ticker.C: + fmt.Fprintf(w, ": ping\n\n") + flusher.Flush() + } + } +} + +// addSSEClient 注册SSE客户端 +func (s *Server) addSSEClient(client *sseClient) { + s.mu.Lock() + defer s.mu.Unlock() + s.sseClients[client.id] = client +} + +// removeSSEClient 移除SSE客户端 +func (s *Server) removeSSEClient(id string) { + s.mu.Lock() + defer s.mu.Unlock() + if client, exists := s.sseClients[id]; exists { + close(client.send) + delete(s.sseClients, id) + } +} + +// handleMessage 处理MCP消息 +func (s *Server) handleMessage(ctx context.Context, msg *Message) *Message { + // 检查是否是通知（notification）- 通知没有id字段，不需要响应 + isNotification := msg.ID.Value() == nil || msg.ID.String() == "" + + // 如果不是通知且ID为空，生成新的UUID + if !isNotification && msg.ID.String() == "" { + msg.ID = MessageID{value: uuid.New().String()} + } + + switch msg.Method { + case "initialize": + return s.handleInitialize(msg) + case "tools/list": + return s.handleListTools(msg) + case "tools/call": + return s.handleCallTool(ctx, msg) + case "prompts/list": + return s.handleListPrompts(msg) + case "prompts/get": + return s.handleGetPrompt(msg) + case "resources/list": + return s.handleListResources(msg) + case "resources/read": + return s.handleReadResource(msg) + case "sampling/request": + return s.handleSamplingRequest(msg) + case "notifications/initialized": + // 通知类型，不需要响应 + s.logger.Debug("收到 initialized 通知") + return nil + case "": + // 空方法名，可能是通知，不返回错误 + if isNotification { + s.logger.Debug("收到无方法名的通知消息") + return nil + } + fallthrough + default: + // 如果是通知，不返回错误响应 + if isNotification { + s.logger.Debug("收到未知通知", zap.String("method", msg.Method)) + return nil + } + // 对于请求，返回方法未找到错误 + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Method not found"}, + } + } +} + +// handleInitialize 处理初始化请求 +func (s *Server) handleInitialize(msg *Message) *Message { + var req InitializeRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + response := InitializeResponse{ + ProtocolVersion: ProtocolVersion, + Capabilities: ServerCapabilities{ + Tools: map[string]interface{}{ + "listChanged": true, + }, + Prompts: map[string]interface{}{ + "listChanged": true, + }, + Resources: map[string]interface{}{ + "subscribe": true, + "listChanged": true, + }, + Sampling: map[string]interface{}{}, + }, + ServerInfo: ServerInfo{ + Name: "CyberStrikeAI", + Version: "1.0.0", + }, + } + + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleListTools 处理列出工具请求 +func (s *Server) handleListTools(msg *Message) *Message { + s.mu.RLock() + tools := make([]Tool, 0, len(s.toolDefs)) + for _, tool := range s.toolDefs { + tools = append(tools, tool) + } + s.mu.RUnlock() + s.logger.Debug("tools/list 请求", zap.Int("返回工具数", len(tools))) + + response := ListToolsResponse{Tools: tools} + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleCallTool 处理工具调用请求 +func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Message { + var req CallToolRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + _, authenticated := authctx.PrincipalFromContext(requestCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(requestCtx, req.Name, req.Arguments); err != nil { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Forbidden", Data: err.Error()}} + } + } else if authenticated { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Tool authorization policy is not configured"}} + } + + executionID := uuid.New().String() + execution := &ToolExecution{ + ID: executionID, + ToolName: req.Name, + Arguments: req.Arguments, + Status: "running", + StartTime: time.Now(), + } + if principal, ok := authctx.PrincipalFromContext(requestCtx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(requestCtx) + + s.mu.Lock() + s.executions[executionID] = execution + // 如果内存中的执行记录超过限制，清理最旧的记录 + s.cleanupOldExecutions() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.mu.RLock() + handler, exists := s.tools[req.Name] + s.mu.RUnlock() + + if !exists { + execution.Status = "failed" + execution.Error = "Tool not found" + now := time.Now() + execution.EndTime = &now + execution.Duration = now.Sub(execution.StartTime) + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + s.mu.Lock() + delete(s.executions, executionID) + s.mu.Unlock() + } + + s.updateStats(req.Name, true) + + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Tool not found"}, + } + } + + baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline(requestCtx) + defer timeoutCancel() + execCtx, runCancel := context.WithCancel(baseCtx) + s.registerRunningCancel(executionID, runCancel) + defer func() { + runCancel() + s.unregisterRunningCancel(executionID) + }() + + s.logger.Info("开始执行工具", + zap.String("toolName", req.Name), + zap.Any("arguments", req.Arguments), + ) + + result, err := handler(execCtx, req.Arguments) + cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err) + now := time.Now() + var failed bool + var finalResult *ToolResult + + s.mu.Lock() + execution.EndTime = &now + execution.Duration = now.Sub(execution.StartTime) + + if err != nil { + st, msg := executionStatusAndMessage(err) + execution.Status = st + execution.Error = msg + failed = st != "cancelled" + } else if result != nil && result.IsError { + if cancelledWithUserNote { + execution.Status = "cancelled" + execution.Error = "" + execution.Result = result + failed = false + } else { + execution.Status = "failed" + if len(result.Content) > 0 { + execution.Error = result.Content[0].Text + } else { + execution.Error = "工具执行返回错误结果" + } + execution.Result = result + failed = true + } + } else { + execution.Status = "completed" + if result == nil { + result = &ToolResult{ + Content: []Content{ + {Type: "text", Text: "工具执行完成，但未返回结果"}, + }, + } + } + execution.Result = result + failed = false + } + + finalResult = execution.Result + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.updateStats(req.Name, failed) + + if s.storage != nil { + s.mu.Lock() + delete(s.executions, executionID) + s.mu.Unlock() + } + + if err != nil { + s.logger.Error("工具执行失败", + zap.String("toolName", req.Name), + zap.Error(err), + ) + + errText := fmt.Sprintf("工具执行失败: %v", err) + if errors.Is(err, context.Canceled) { + errText = "工具执行已手动终止（MCP 监控）。后续编排步骤可继续。" + } + errorResult, _ := json.Marshal(CallToolResponse{ + Content: []Content{ + {Type: "text", Text: errText}, + }, + IsError: true, + }) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: errorResult, + } + } + + if finalResult != nil && finalResult.IsError { + s.logger.Warn("工具执行返回错误结果", + zap.String("toolName", req.Name), + ) + + errorResult, _ := json.Marshal(CallToolResponse{ + Content: finalResult.Content, + IsError: true, + }) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: errorResult, + } + } + + if finalResult == nil { + finalResult = &ToolResult{ + Content: []Content{ + {Type: "text", Text: "工具执行完成，但未返回结果"}, + }, + } + } + + resultJSON, _ := json.Marshal(CallToolResponse{ + Content: finalResult.Content, + IsError: false, + }) + + s.logger.Info("工具执行完成", + zap.String("toolName", req.Name), + zap.Bool("isError", finalResult.IsError), + ) + + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: resultJSON, + } +} + +// updateStats 更新统计信息 +func (s *Server) updateStats(toolName string, failed bool) { + now := time.Now() + if s.storage != nil { + totalCalls := 1 + successCalls := 0 + failedCalls := 0 + if failed { + failedCalls = 1 + } else { + successCalls = 1 + } + if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil { + s.logger.Warn("保存统计信息到数据库失败", zap.Error(err)) + } + return + } + + s.mu.Lock() + defer s.mu.Unlock() + + if s.stats[toolName] == nil { + s.stats[toolName] = &ToolStats{ + ToolName: toolName, + } + } + + stats := s.stats[toolName] + stats.TotalCalls++ + stats.LastCallTime = &now + + if failed { + stats.FailedCalls++ + } else { + stats.SuccessCalls++ + } +} + +// GetExecution 获取执行记录（先从内存查找，再从数据库查找） +func (s *Server) GetExecution(id string) (*ToolExecution, bool) { + if s.executionService != nil { + if snap, err := s.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil { + return snap.Execution, true + } + } + s.mu.RLock() + exec, exists := s.executions[id] + s.mu.RUnlock() + + if exists { + return exec, true + } + + if s.storage != nil { + exec, err := s.storage.GetToolExecution(id) + if err == nil { + return exec, true + } + } + + return nil, false +} + +// loadHistoricalData 从数据库加载历史数据 +func (s *Server) loadHistoricalData() { + if s.storage == nil { + return + } + + // 加载历史执行记录（最近1000条） + executions, err := s.storage.LoadToolExecutions() + if err != nil { + s.logger.Warn("加载历史执行记录失败", zap.Error(err)) + } else { + s.mu.Lock() + for _, exec := range executions { + // 只加载最近 maxExecutionsInMemory 条，避免内存占用过大 + if len(s.executions) < s.maxExecutionsInMemory { + s.executions[exec.ID] = exec + } else { + break + } + } + s.mu.Unlock() + s.logger.Info("加载历史执行记录", zap.Int("count", len(executions))) + } + + // 加载历史统计信息 + stats, err := s.storage.LoadToolStats() + if err != nil { + s.logger.Warn("加载历史统计信息失败", zap.Error(err)) + } else { + s.mu.Lock() + for k, v := range stats { + s.stats[k] = v + } + s.mu.Unlock() + s.logger.Info("加载历史统计信息", zap.Int("count", len(stats))) + } +} + +// GetAllExecutions 获取所有执行记录（合并内存和数据库） +func (s *Server) GetAllExecutions() []*ToolExecution { + if s.storage != nil { + dbExecutions, err := s.storage.LoadToolExecutions() + if err == nil { + execMap := make(map[string]*ToolExecution) + for _, exec := range dbExecutions { + if _, exists := execMap[exec.ID]; !exists { + execMap[exec.ID] = exec + } + } + + s.mu.RLock() + for id, exec := range s.executions { + if _, exists := execMap[id]; !exists { + execMap[id] = exec + } + } + s.mu.RUnlock() + + result := make([]*ToolExecution, 0, len(execMap)) + for _, exec := range execMap { + result = append(result, exec) + } + return result + } else { + s.logger.Warn("从数据库加载执行记录失败", zap.Error(err)) + } + } + + s.mu.RLock() + defer s.mu.RUnlock() + + memExecutions := make([]*ToolExecution, 0, len(s.executions)) + for _, exec := range s.executions { + memExecutions = append(memExecutions, exec) + } + return memExecutions +} + +// GetStats 获取统计信息（合并内存和数据库） +func (s *Server) GetStats() map[string]*ToolStats { + if s.storage != nil { + dbStats, err := s.storage.LoadToolStats() + if err == nil { + return dbStats + } + s.logger.Warn("从数据库加载统计信息失败", zap.Error(err)) + } + + s.mu.RLock() + defer s.mu.RUnlock() + + memStats := make(map[string]*ToolStats) + for k, v := range s.stats { + statCopy := *v + memStats[k] = &statCopy + } + + return memStats +} + +// GetAllTools 获取所有已注册的工具（用于Agent动态获取工具列表） +func (s *Server) GetAllTools() []Tool { + s.mu.RLock() + defer s.mu.RUnlock() + + tools := make([]Tool, 0, len(s.toolDefs)) + for _, tool := range s.toolDefs { + tools = append(tools, tool) + } + return tools +} + +// CallTool 直接调用工具（用于内部调用） +func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { + if s.executionService == nil { + s.executionService = NewExecutionService(s.storage, s.logger) + s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes) + s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir) + } + var ownerUserID string + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + ownerUserID = principal.UserID + } + handle, err := s.executionService.Submit(ctx, ExecutionRequest{ + ToolName: toolName, + Arguments: args, + ConversationID: MCPConversationIDFromContext(ctx), + OwnerUserID: ownerUserID, + Run: func(runCtx context.Context) (*ToolResult, error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + handler, exists := s.tools[toolName] + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("tool authorization denied: %w", err) + } + } else if authenticated { + return nil, errors.New("tool authorization policy is not configured") + } + if !exists { + return nil, fmt.Errorf("工具 %s 未找到", toolName) + } + return handler(runCtx, args) + }, + OnDone: func(exec *ToolExecution) { + failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled + s.updateStats(toolName, failed) + }, + }) + if err != nil { + return nil, "", err + } + + s.mu.RLock() + waitTimeout := s.toolWaitTimeout + s.mu.RUnlock() + if isExecutionControlTool(toolName) { + waitTimeout = 0 + } + snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout) + if errors.Is(waitErr, ErrExecutionWaitTimeout) { + return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil + } + if waitErr != nil { + return nil, handle.ID, waitErr + } + if snapshot == nil || snapshot.Execution == nil { + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回执行快照"}}, IsError: true}, handle.ID, nil + } + if snapshot.Execution.Result != nil { + return snapshot.Execution.Result, handle.ID, nil + } + if snapshot.Execution.Error != "" { + return nil, handle.ID, errors.New(snapshot.Execution.Error) + } + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}, IsError: false}, handle.ID, nil +} + +func internalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult { + execID := "" + status := ToolExecutionStatusRunning + toolName := "" + elapsed := time.Duration(0) + if snapshot != nil && snapshot.Execution != nil { + execID = snapshot.Execution.ID + status = snapshot.Execution.Status + toolName = snapshot.Execution.ToolName + elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second) + } + waitText := "unbounded" + if waitTimeout > 0 { + waitText = waitTimeout.Round(time.Second).String() + } + msg := fmt.Sprintf(`工具已提交到后台执行，但本次等待已到达上限。 + +execution_id: %s +tool: %s +status: %s +wait_timeout: %s +elapsed: %s + +你可以继续推理、改用其他工具，或调用 wait_tool_execution 继续等待该 execution_id；也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed) + return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true} +} + +func isExecutionControlTool(toolName string) bool { + switch strings.TrimSpace(toolName) { + case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution, builtin.ToolCancelToolExecution: + return true + default: + return false + } +} + +// BeginToolExecution 创建 running 状态的执行记录，供 Eino 等非 CallTool 路径在工具开始时落库。 +func (s *Server) BeginToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string { + if s == nil { + return "" + } + if args == nil { + args = map[string]interface{}{} + } + executionID := uuid.New().String() + execution := &ToolExecution{ + ID: executionID, + ToolName: toolName, + Arguments: args, + Status: "running", + StartTime: time.Now(), + } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(ctx) + + s.mu.Lock() + s.executions[executionID] = execution + s.cleanupOldExecutions() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + return executionID +} + +// FinishToolExecution 完成先前 BeginToolExecution 创建的记录；executionID 为空时等同 RecordCompletedToolInvocation。 +func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { + if s == nil { + return "" + } + if args == nil { + args = map[string]interface{}{} + } + id := strings.TrimSpace(executionID) + if id == "" { + id = uuid.New().String() + } + + now := time.Now() + failed := invokeErr != nil + var finalResult *ToolResult + + s.mu.Lock() + maxBytes := s.toolResultMaxBytes + spillRoot := s.spillRootDir + exec, inMem := s.executions[id] + if !inMem || exec == nil { + exec = &ToolExecution{ + ID: id, + ToolName: toolName, + Arguments: args, + StartTime: now, + } + s.executions[id] = exec + } else if toolName != "" { + exec.ToolName = toolName + } + if len(args) > 0 { + exec.Arguments = args + } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + exec.OwnerUserID = principal.UserID + } + if conversationID := MCPConversationIDFromContext(ctx); conversationID != "" { + exec.ConversationID = conversationID + } + exec.EndTime = &now + if exec.StartTime.IsZero() { + exec.StartTime = now + } + exec.Duration = now.Sub(exec.StartTime) + + spill := ToolResultSpillConfig{ + RootDir: spillRoot, + ProjectID: MCPProjectIDFromContext(ctx), + ConversationID: exec.ConversationID, + ExecutionID: id, + } + if failed { + st, msg := executionStatusAndMessage(invokeErr) + exec.Status = st + exec.Error = msg + if strings.TrimSpace(resultText) != "" { + finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}} + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) + exec.Result = finalResult + } + } else { + exec.Status = "completed" + text := resultText + if strings.TrimSpace(text) == "" { + text = "（无输出）" + } + finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}} + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) + exec.Result = finalResult + } + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(exec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.updateStats(exec.ToolName, failed) + + if s.storage != nil { + s.mu.Lock() + delete(s.executions, id) + s.mu.Unlock() + } + return id +} + +// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution. +// The final Result remains authoritative and is written only when the tool finishes. +func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) { + if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" { + return + } + id := strings.TrimSpace(executionID) + if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) { + return + } + now := time.Now() + s.mu.Lock() + exec := s.executions[id] + if exec != nil { + appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now) + } + s.mu.Unlock() +} + +// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储（格式与 CallTool 结束后一致）， +// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景；返回 executionId 供助手消息 mcpExecutionIds 关联。 +func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { + return s.FinishToolExecution(ctx, "", toolName, args, resultText, invokeErr) +} + +// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文（如 reduction 后的 persisted-output）。 +func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResult) error { + if s == nil { + return nil + } + executionID = strings.TrimSpace(executionID) + if executionID == "" || result == nil { + return nil + } + s.mu.Lock() + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ExecutionID: executionID, + } + if exec, ok := s.executions[executionID]; ok && exec != nil { + spill.ConversationID = exec.ConversationID + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) + exec.Result = result + } else { + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) + } + s.mu.Unlock() + if s.storage != nil { + return s.storage.UpdateToolExecutionResult(executionID, result) + } + return nil +} + +// cleanupOldExecutions 清理旧的执行记录，防止内存无限增长 +func (s *Server) cleanupOldExecutions() { + if len(s.executions) <= s.maxExecutionsInMemory { + return + } + + // 按开始时间排序，找出最旧的记录 + type execWithTime struct { + id string + startTime time.Time + } + execs := make([]execWithTime, 0, len(s.executions)) + for id, exec := range s.executions { + execs = append(execs, execWithTime{ + id: id, + startTime: exec.StartTime, + }) + } + + // 使用 sort 包进行高效排序（最旧的在前） + sort.Slice(execs, func(i, j int) bool { + return execs[i].startTime.Before(execs[j].startTime) + }) + + // 删除最旧的记录，保留 maxExecutionsInMemory 条 + toDelete := len(s.executions) - s.maxExecutionsInMemory + for i := 0; i < toDelete; i++ { + delete(s.executions, execs[i].id) + } + + s.logger.Debug("清理旧的执行记录", + zap.Int("before", len(execs)), + zap.Int("after", len(s.executions)), + zap.Int("deleted", toDelete), + ) +} + +func (s *Server) registerRunningCancel(id string, cancel context.CancelFunc) { + s.runningCancelsMu.Lock() + s.runningCancels[id] = cancel + s.runningCancelsMu.Unlock() +} + +func (s *Server) unregisterRunningCancel(id string) { + s.runningCancelsMu.Lock() + delete(s.runningCancels, id) + s.runningCancelsMu.Unlock() +} + +// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino +// filesystem execute, participate in cancel_tool_execution by execution_id. +func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) { + id = strings.TrimSpace(id) + if s == nil || id == "" || cancel == nil { + return + } + s.registerRunningCancel(id, cancel) +} + +func (s *Server) UnregisterToolExecutionCancel(id string) { + id = strings.TrimSpace(id) + if s == nil || id == "" { + return + } + s.unregisterRunningCancel(id) +} + +func (s *Server) readAbortUserNote(id string) string { + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if s.abortUserNotes == nil { + return "" + } + return s.abortUserNotes[id] +} + +func (s *Server) takeAbortUserNote(id string) string { + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if s.abortUserNotes == nil { + return "" + } + n := s.abortUserNotes[id] + delete(s.abortUserNotes, id) + return n +} + +// applyAbortUserNoteToCancelledToolResult 监控页「终止并填写说明」时合并「工具已输出 + 用户说明」交给模型。 +// exec 等工具会把失败写在 *ToolResult 里并返回 err==nil，若仅在 err!=nil 时合并会漏掉说明，甚至误 clear 掉 note。 +func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(s.readAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + _ = s.takeAbortUserNote(executionID) + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + *err = nil + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + return true +} + +// CancelToolExecutionWithNote 取消内部工具；note 非空时与工具已返回文本合并后交给上层模型。 +func (s *Server) CancelToolExecutionWithNote(id string, note string) bool { + if s.executionService != nil && s.executionService.Cancel(id, note) { + return true + } + s.runningCancelsMu.Lock() + cancel, ok := s.runningCancels[id] + if !ok || cancel == nil { + s.runningCancelsMu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + if s.abortUserNotes == nil { + s.abortUserNotes = make(map[string]string) + } + s.abortUserNotes[id] = strings.TrimSpace(note) + } + s.runningCancelsMu.Unlock() + cancel() + return true +} + +// CancelToolExecution 取消正在执行的内部工具调用（无用户说明）。 +func (s *Server) CancelToolExecution(id string) bool { + return s.CancelToolExecutionWithNote(id, "") +} + +// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的 executionId 快照。 +func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} { + if s == nil { + return nil + } + out := make(map[string]struct{}) + if s.executionService != nil { + for id := range s.executionService.ActiveRunningExecutionIDs() { + out[id] = struct{}{} + } + } + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if len(s.runningCancels) == 0 && len(out) == 0 { + return nil + } + for id := range s.runningCancels { + out[id] = struct{}{} + } + return out +} + +// initDefaultPrompts 初始化默认提示词模板 +func (s *Server) initDefaultPrompts() { + s.mu.Lock() + defer s.mu.Unlock() + + // 网络安全测试提示词 + s.prompts["security_scan"] = &Prompt{ + Name: "security_scan", + Description: "生成网络安全扫描任务的提示词", + Arguments: []PromptArgument{ + {Name: "target", Description: "扫描目标（IP地址或域名）", Required: true}, + {Name: "scan_type", Description: "扫描类型（port, vuln, web等）", Required: false}, + }, + } + + // 渗透测试提示词 + s.prompts["penetration_test"] = &Prompt{ + Name: "penetration_test", + Description: "生成渗透测试任务的提示词", + Arguments: []PromptArgument{ + {Name: "target", Description: "测试目标", Required: true}, + {Name: "scope", Description: "测试范围", Required: false}, + }, + } +} + +// initDefaultResources 初始化默认资源 +// 注意：工具资源现在在 RegisterTool 时自动创建，此函数保留用于其他非工具资源 +func (s *Server) initDefaultResources() { + // 工具资源已改为在 RegisterTool 时自动创建，无需在此硬编码 +} + +// handleListPrompts 处理列出提示词请求 +func (s *Server) handleListPrompts(msg *Message) *Message { + s.mu.RLock() + prompts := make([]Prompt, 0, len(s.prompts)) + for _, prompt := range s.prompts { + prompts = append(prompts, *prompt) + } + s.mu.RUnlock() + + response := ListPromptsResponse{ + Prompts: prompts, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleGetPrompt 处理获取提示词请求 +func (s *Server) handleGetPrompt(msg *Message) *Message { + var req GetPromptRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + s.mu.RLock() + prompt, exists := s.prompts[req.Name] + s.mu.RUnlock() + + if !exists { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Prompt not found"}, + } + } + + // 根据提示词名称生成消息 + messages := s.generatePromptMessages(prompt, req.Arguments) + + response := GetPromptResponse{ + Messages: messages, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// generatePromptMessages 生成提示词消息 +func (s *Server) generatePromptMessages(prompt *Prompt, args map[string]interface{}) []PromptMessage { + messages := []PromptMessage{} + + switch prompt.Name { + case "security_scan": + target, _ := args["target"].(string) + scanType, _ := args["scan_type"].(string) + if scanType == "" { + scanType = "comprehensive" + } + + content := fmt.Sprintf(`请对目标 %s 执行%s安全扫描。包括： +1. 端口扫描和服务识别 +2. 漏洞检测 +3. Web应用安全测试 +4. 生成详细的安全报告`, target, scanType) + + messages = append(messages, PromptMessage{ + Role: "user", + Content: content, + }) + + case "penetration_test": + target, _ := args["target"].(string) + scope, _ := args["scope"].(string) + + content := fmt.Sprintf(`请对目标 %s 执行渗透测试。`, target) + if scope != "" { + content += fmt.Sprintf("测试范围：%s", scope) + } + content += "\n请按照OWASP Top 10进行全面的安全测试。" + + messages = append(messages, PromptMessage{ + Role: "user", + Content: content, + }) + + default: + messages = append(messages, PromptMessage{ + Role: "user", + Content: "请执行安全测试任务", + }) + } + + return messages +} + +// handleListResources 处理列出资源请求 +func (s *Server) handleListResources(msg *Message) *Message { + s.mu.RLock() + resources := make([]Resource, 0, len(s.resources)) + for _, resource := range s.resources { + resources = append(resources, *resource) + } + s.mu.RUnlock() + + response := ListResourcesResponse{ + Resources: resources, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleReadResource 处理读取资源请求 +func (s *Server) handleReadResource(msg *Message) *Message { + var req ReadResourceRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + s.mu.RLock() + resource, exists := s.resources[req.URI] + s.mu.RUnlock() + + if !exists { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Resource not found"}, + } + } + + // 生成资源内容 + content := s.generateResourceContent(resource) + + response := ReadResourceResponse{ + Contents: []ResourceContent{content}, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// generateResourceContent 生成资源内容 +func (s *Server) generateResourceContent(resource *Resource) ResourceContent { + content := ResourceContent{ + URI: resource.URI, + MimeType: resource.MimeType, + } + + // 如果是工具资源，生成详细文档 + if strings.HasPrefix(resource.URI, "tool://") { + toolName := strings.TrimPrefix(resource.URI, "tool://") + content.Text = s.generateToolDocumentation(toolName, resource) + } else { + // 其他资源使用描述或默认内容 + content.Text = resource.Description + } + + return content +} + +// generateToolDocumentation 生成工具文档 +// 注意：硬编码的工具文档已移除，现在只使用工具定义中的信息 +func (s *Server) generateToolDocumentation(toolName string, resource *Resource) string { + // 获取工具定义以获取更详细的信息 + s.mu.RLock() + tool, hasTool := s.toolDefs[toolName] + s.mu.RUnlock() + + // 使用工具定义中的描述信息 + if hasTool { + doc := fmt.Sprintf("%s\n\n", resource.Description) + if tool.InputSchema != nil { + if props, ok := tool.InputSchema["properties"].(map[string]interface{}); ok { + doc += "参数说明：\n" + for paramName, paramInfo := range props { + if paramMap, ok := paramInfo.(map[string]interface{}); ok { + if desc, ok := paramMap["description"].(string); ok { + doc += fmt.Sprintf("- %s: %s\n", paramName, desc) + } + } + } + } + } + return doc + } + return resource.Description +} + +// handleSamplingRequest 处理采样请求 +func (s *Server) handleSamplingRequest(msg *Message) *Message { + var req SamplingRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + // 注意：采样功能通常需要连接到实际的LLM服务 + // 这里返回一个占位符响应，实际实现需要集成LLM API + s.logger.Warn("Sampling request received but not fully implemented", + zap.Any("request", req), + ) + + response := SamplingResponse{ + Content: []SamplingContent{ + { + Type: "text", + Text: "采样功能需要配置LLM服务。请使用Agent Loop API进行AI对话。", + }, + }, + StopReason: "length", + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// RegisterPrompt 注册提示词模板 +func (s *Server) RegisterPrompt(prompt *Prompt) { + s.mu.Lock() + defer s.mu.Unlock() + s.prompts[prompt.Name] = prompt +} + +// RegisterResource 注册资源 +func (s *Server) RegisterResource(resource *Resource) { + s.mu.Lock() + defer s.mu.Unlock() + s.resources[resource.URI] = resource +} + +// HandleStdio 处理标准输入输出（用于 stdio 传输模式） +// MCP 协议使用换行分隔的 JSON-RPC 消息；管道下需每次写入后 Flush，否则客户端会读不到响应 +func (s *Server) HandleStdio() error { + decoder := json.NewDecoder(os.Stdin) + stdout := bufio.NewWriter(os.Stdout) + encoder := json.NewEncoder(stdout) + // 注意：不设置缩进，MCP 协议期望紧凑的 JSON 格式 + + for { + var msg Message + if err := decoder.Decode(&msg); err != nil { + if err == io.EOF { + break + } + // 日志输出到 stderr，避免干扰 stdout 的 JSON-RPC 通信 + s.logger.Error("读取消息失败", zap.Error(err)) + // 发送错误响应 + errorMsg := Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32700, Message: "Parse error", Data: err.Error()}, + } + if err := encoder.Encode(errorMsg); err != nil { + return fmt.Errorf("发送错误响应失败: %w", err) + } + if err := stdout.Flush(); err != nil { + return fmt.Errorf("刷新 stdout 失败: %w", err) + } + continue + } + + // 处理消息 + response := s.handleMessage(context.Background(), &msg) + + // 如果是通知（response 为 nil），不需要发送响应 + if response == nil { + continue + } + + // 发送响应 + if err := encoder.Encode(response); err != nil { + return fmt.Errorf("发送响应失败: %w", err) + } + if err := stdout.Flush(); err != nil { + return fmt.Errorf("刷新 stdout 失败: %w", err) + } + } + + return nil +} + +// sendError 发送错误响应 +func (s *Server) sendError(w http.ResponseWriter, id interface{}, code int, message, data string) { + var msgID MessageID + if id != nil { + msgID = MessageID{value: id} + } + response := Message{ + ID: msgID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: code, Message: message, Data: data}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(response) +} diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go new file mode 100644 index 00000000..bea489a3 --- /dev/null +++ b/internal/mcp/server_authorization_test.go @@ -0,0 +1,231 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + + "go.uber.org/zap" +) + +func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) { + server := NewServer(zap.NewNop()) + server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil + }) + server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error { + if _, ok := authctx.PrincipalFromContext(ctx); !ok { + return errors.New("principal required") + } + return nil + }) + _, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil) + if err == nil { + t.Fatal("tool call without principal was allowed") + } + if deniedExecutionID == "" { + t.Fatal("denied tool call should still return an execution id") + } + deniedExecution, ok := server.GetExecution(deniedExecutionID) + if !ok || deniedExecution == nil { + t.Fatalf("missing denied execution %q", deniedExecutionID) + } + if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") { + t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution) + } + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true})) + _, executionID, err := server.CallTool(ctx, "echo", nil) + if err != nil { + t.Fatal(err) + } + execution, ok := server.GetExecution(executionID) + if !ok || execution.OwnerUserID != "u1" { + t.Fatalf("execution owner = %#v, want u1", execution) + } +} + +func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + started := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + close(started) + select { + case <-release: + return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + + callCtx, callCancel := context.WithCancel(context.Background()) + result, executionID, err := server.CallTool(callCtx, "slow", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if executionID == "" || result == nil || !result.IsError { + t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID) + } + if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") { + t.Fatalf("timeout result missing execution guidance: %q", text) + } + select { + case <-started: + default: + t.Fatal("internal worker did not start") + } + callCancel() + close(release) + + snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second) + if err != nil { + t.Fatalf("wait internal execution: %v", err) + } + if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("snapshot = %#v, want completed", snapshot) + } + if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" { + t.Fatalf("result = %q, want internal done", got) + } +} + +func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + select { + case <-release: + return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + RegisterExecutionControlTools(server, nil) + + result, executionID, err := server.CallTool(context.Background(), "slow", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID) + } + + done := make(chan *ToolResult, 1) + errCh := make(chan error, 1) + go func() { + waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 1, + }) + if waitErr != nil { + errCh <- waitErr + return + } + done <- waitResult + }() + + select { + case <-done: + t.Fatal("wait_tool_execution returned before target execution completed") + case err := <-errCh: + t.Fatalf("wait_tool_execution errored before release: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(release) + + select { + case err := <-errCh: + t.Fatalf("wait_tool_execution returned error: %v", err) + case waitResult := <-done: + if waitResult == nil || waitResult.IsError { + t.Fatalf("expected successful wait result, got %#v", waitResult) + } + if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) { + t.Fatalf("wait result missing completed target: %s", body) + } + case <-time.After(time.Second): + t.Fatal("wait_tool_execution did not return after target completion") + } +} + +func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + <-release + return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil + }) + RegisterExecutionControlTools(server, nil) + + result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID) + } + + waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 0.01, + }) + if err != nil { + t.Fatalf("wait_tool_execution returned error: %v", err) + } + if waitResult == nil { + t.Fatal("missing wait result") + } + if waitResult.IsError { + t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult) + } + body := ToolResultPlainText(waitResult) + if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") { + t.Fatalf("wait timeout body missing running status/guidance: %s", body) + } + close(release) +} + +func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) { + server := NewServer(zap.NewNop()) + RegisterExecutionControlTools(server, nil) + + executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"}) + if executionID == "" { + t.Fatal("missing execution id") + } + server.AppendToolExecutionPartialOutput(executionID, "first\n") + server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32)) + + result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "partial_output_max_bytes": 8, + }) + if err != nil { + t.Fatalf("get_tool_execution: %v", err) + } + body := ToolResultPlainText(result) + if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) { + t.Fatalf("missing bounded partial output: %s", body) + } + if !strings.Contains(body, `"partial_output_bytes": 38`) { + t.Fatalf("missing partial byte count: %s", body) + } + + result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "include_partial_output": false, + }) + if err != nil { + t.Fatalf("get_tool_execution without partial: %v", err) + } + if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") { + t.Fatalf("partial output should be omitted: %s", body) + } +} diff --git a/internal/mcp/tool_result_guard.go b/internal/mcp/tool_result_guard.go new file mode 100644 index 00000000..851db4df --- /dev/null +++ b/internal/mcp/tool_result_guard.go @@ -0,0 +1,65 @@ +package mcp + +import "cyberstrike-ai/internal/tooloutput" + +const DefaultToolResultMaxBytes = 12000 + +// ToolResultSpillConfig controls where oversized tool results are written on disk +// before the in-memory/DB/agent-facing payload is truncated. +type ToolResultSpillConfig struct { + RootDir string + ProjectID string + ConversationID string + ExecutionID string +} + +// NormalizeToolResultForStorage returns the canonical result used by both the +// agent-facing response and monitor persistence. When maxBytes is exceeded the +// full text is spilled under the reduction cache tree and replaced with a +// notice that includes the file path. +func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult { + return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{}) +} + +// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with +// an explicit spill location (conversation/execution scoped). +func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult { + if result == nil { + return nil + } + out := cloneToolResult(result) + if maxBytes <= 0 { + return out + } + + total := 0 + for _, c := range out.Content { + if c.Type == "text" { + total += len(c.Text) + } + } + if total <= maxBytes { + return out + } + + full := ToolResultPlainText(out) + bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{ + RootDir: spill.RootDir, + ProjectID: spill.ProjectID, + ConversationID: spill.ConversationID, + ExecutionID: spill.ExecutionID, + }) + out.Content = []Content{{Type: "text", Text: bound}} + return out +} + +func cloneToolResult(in *ToolResult) *ToolResult { + if in == nil { + return nil + } + out := *in + if in.Content != nil { + out.Content = append([]Content(nil), in.Content...) + } + return &out +} diff --git a/internal/mcp/tool_result_guard_test.go b/internal/mcp/tool_result_guard_test.go new file mode 100644 index 00000000..fd7c8d55 --- /dev/null +++ b/internal/mcp/tool_result_guard_test.go @@ -0,0 +1,158 @@ +package mcp + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +type inMemoryMonitorStorage struct { + executions map[string]*ToolExecution +} + +func newInMemoryMonitorStorage() *inMemoryMonitorStorage { + return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}} +} + +func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error { + if exec != nil { + s.executions[exec.ID] = cloneToolExecution(exec) + } + return nil +} + +func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error { + exec := s.executions[id] + if exec == nil { + exec = &ToolExecution{ID: id} + s.executions[id] = exec + } + exec.Result = cloneToolResult(result) + return nil +} + +func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) { + out := make([]*ToolExecution, 0, len(s.executions)) + for _, exec := range s.executions { + out = append(out, cloneToolExecution(exec)) + } + return out, nil +} + +func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) { + if exec := s.executions[id]; exec != nil { + return cloneToolExecution(exec), nil + } + return nil, nil +} + +func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil } + +func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) { + return map[string]*ToolStats{}, nil +} + +func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error { + return nil +} + +func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) { + storage := newInMemoryMonitorStorage() + server := NewServerWithStorage(zap.NewNop(), storage) + server.ConfigureToolWaitTimeoutSeconds(0) + server.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + server.ConfigureToolResultSpillRoot(spillRoot) + server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil + }) + + ctx := WithMCPConversationID(context.Background(), "conv-spill") + result, executionID, err := server.CallTool(ctx, "big", nil) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if executionID == "" { + t.Fatal("missing execution id") + } + returned := ToolResultPlainText(result) + if !strings.Contains(returned, "") || !strings.Contains(returned, "Full output saved to:") { + t.Fatalf("returned result was not spilled: %q", returned) + } + if len(returned) > 400 { + t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned) + } + + spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID) + abs, err := filepath.Abs(spillPath) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(returned, abs) { + t.Fatalf("missing spill path %q in %q", abs, returned) + } + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill file: %v", err) + } + if string(body) != strings.Repeat("x", 800) { + t.Fatalf("spill body mismatch: len=%d", len(body)) + } + + inMem, ok := server.GetExecution(executionID) + if !ok || inMem == nil || inMem.Result == nil { + t.Fatalf("missing in-memory execution: %#v", inMem) + } + stored := storage.executions[executionID] + if stored == nil || stored.Result == nil { + t.Fatalf("missing stored execution: %#v", stored) + } + if ToolResultPlainText(inMem.Result) != returned { + t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned) + } + if ToolResultPlainText(stored.Result) != returned { + t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned) + } +} + +func TestExecutionServiceStoresGuardedResult(t *testing.T) { + service := NewExecutionService(nil, zap.NewNop()) + service.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + service.ConfigureToolResultSpillRoot(spillRoot) + handle, err := service.Submit(context.Background(), ExecutionRequest{ + ToolName: "big", + ConversationID: "svc-conv", + Run: func(context.Context) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil + }, + }) + if err != nil { + t.Fatalf("Submit: %v", err) + } + snap, err := service.Wait(context.Background(), handle.ID, time.Second) + if err != nil { + t.Fatalf("Wait: %v", err) + } + got := ToolResultPlainText(snap.Execution.Result) + if !strings.Contains(got, "") { + t.Fatalf("service result was not spilled: %q", got) + } + if len(got) > 400 { + t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got) + } + path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID) + abs, _ := filepath.Abs(path) + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill: %v", err) + } + if string(body) != strings.Repeat("a", 800) { + t.Fatalf("unexpected spill body len=%d", len(body)) + } +} diff --git a/internal/mcp/types.go b/internal/mcp/types.go new file mode 100644 index 00000000..6922e047 --- /dev/null +++ b/internal/mcp/types.go @@ -0,0 +1,338 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" +) + +// ExternalMCPClient 外部 MCP 客户端接口（由 client_sdk.go 基于官方 SDK 实现） +type ExternalMCPClient interface { + Initialize(ctx context.Context) error + ListTools(ctx context.Context) ([]Tool, error) + CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) + Close() error + IsConnected() bool + GetStatus() string +} + +// MCP消息类型 +const ( + MessageTypeRequest = "request" + MessageTypeResponse = "response" + MessageTypeError = "error" + MessageTypeNotify = "notify" +) + +// MCP协议版本 +const ProtocolVersion = "2024-11-05" + +// MessageID 表示JSON-RPC 2.0的id字段，可以是字符串、数字或null +type MessageID struct { + value interface{} +} + +// UnmarshalJSON 自定义反序列化，支持字符串、数字和null +func (m *MessageID) UnmarshalJSON(data []byte) error { + // 尝试解析为null + if string(data) == "null" { + m.value = nil + return nil + } + + // 尝试解析为字符串 + var str string + if err := json.Unmarshal(data, &str); err == nil { + m.value = str + return nil + } + + // 尝试解析为数字 + var num json.Number + if err := json.Unmarshal(data, &num); err == nil { + m.value = num + return nil + } + + return fmt.Errorf("invalid id type") +} + +// MarshalJSON 自定义序列化 +func (m MessageID) MarshalJSON() ([]byte, error) { + if m.value == nil { + return []byte("null"), nil + } + return json.Marshal(m.value) +} + +// String 返回字符串表示 +func (m MessageID) String() string { + if m.value == nil { + return "" + } + return fmt.Sprintf("%v", m.value) +} + +// Value 返回原始值 +func (m MessageID) Value() interface{} { + return m.value +} + +// Message 表示MCP消息（符合JSON-RPC 2.0规范） +type Message struct { + ID MessageID `json:"id,omitempty"` + Type string `json:"-"` // 内部使用，不序列化到JSON + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *Error `json:"error,omitempty"` + Version string `json:"jsonrpc,omitempty"` // JSON-RPC 2.0 版本标识 +} + +// Error 表示MCP错误 +type Error struct { + Code int `json:"code"` + Message string `json:"message"` + Data interface{} `json:"data,omitempty"` +} + +// Tool 表示MCP工具定义 +type Tool struct { + Name string `json:"name"` + Description string `json:"description"` // 详细描述 + ShortDescription string `json:"shortDescription,omitempty"` // 简短描述（用于工具列表，减少token消耗） + InputSchema map[string]interface{} `json:"inputSchema"` +} + +// ToolCall 表示工具调用 +type ToolCall struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` +} + +// ToolResult 表示工具执行结果 +type ToolResult struct { + Content []Content `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// Content 表示内容 +type Content struct { + Type string `json:"type"` + Text string `json:"text"` +} + +// InitializeRequest 初始化请求 +type InitializeRequest struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]interface{} `json:"capabilities"` + ClientInfo ClientInfo `json:"clientInfo"` +} + +// ClientInfo 客户端信息 +type ClientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// InitializeResponse 初始化响应 +type InitializeResponse struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities ServerCapabilities `json:"capabilities"` + ServerInfo ServerInfo `json:"serverInfo"` +} + +// ServerCapabilities 服务器能力 +type ServerCapabilities struct { + Tools map[string]interface{} `json:"tools,omitempty"` + Prompts map[string]interface{} `json:"prompts,omitempty"` + Resources map[string]interface{} `json:"resources,omitempty"` + Sampling map[string]interface{} `json:"sampling,omitempty"` +} + +// ServerInfo 服务器信息 +type ServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// ListToolsRequest 列出工具请求 +type ListToolsRequest struct{} + +// ListToolsResponse 列出工具响应 +type ListToolsResponse struct { + Tools []Tool `json:"tools"` +} + +// ListPromptsResponse 列出提示词响应 +type ListPromptsResponse struct { + Prompts []Prompt `json:"prompts"` +} + +// ListResourcesResponse 列出资源响应 +type ListResourcesResponse struct { + Resources []Resource `json:"resources"` +} + +// CallToolRequest 调用工具请求 +type CallToolRequest struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` +} + +// CallToolResponse 调用工具响应 +type CallToolResponse struct { + Content []Content `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// ToolExecution 工具执行记录 +type ToolExecution struct { + ID string `json:"id"` + ToolName string `json:"toolName"` + Arguments map[string]interface{} `json:"arguments"` + Status string `json:"status"` // pending, running, completed, failed, cancelled + Result *ToolResult `json:"result,omitempty"` + Error string `json:"error,omitempty"` + StartTime time.Time `json:"startTime"` + EndTime *time.Time `json:"endTime,omitempty"` + Duration time.Duration `json:"duration,omitempty"` + // PartialOutput is a bounded tail preview of output produced by a running tool. + // It is intentionally separate from Result, which remains the final canonical tool result. + PartialOutput string `json:"partialOutput,omitempty"` + PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"` + PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"` + PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"` + // ConversationID 仅 API 展示用（进行中的 Agent 任务），不写入 tool_executions 表。 + ConversationID string `json:"conversationId,omitempty"` + OwnerUserID string `json:"-"` +} + +// ToolStats 工具统计信息 +type ToolStats struct { + ToolName string `json:"toolName"` + TotalCalls int `json:"totalCalls"` + SuccessCalls int `json:"successCalls"` + FailedCalls int `json:"failedCalls"` + LastCallTime *time.Time `json:"lastCallTime,omitempty"` +} + +// Prompt 提示词模板 +type Prompt struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Arguments []PromptArgument `json:"arguments,omitempty"` +} + +// PromptArgument 提示词参数 +type PromptArgument struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Required bool `json:"required,omitempty"` +} + +// GetPromptRequest 获取提示词请求 +type GetPromptRequest struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments,omitempty"` +} + +// GetPromptResponse 获取提示词响应 +type GetPromptResponse struct { + Messages []PromptMessage `json:"messages"` +} + +// PromptMessage 提示词消息 +type PromptMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// Resource 资源 +type Resource struct { + URI string `json:"uri"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + MimeType string `json:"mimeType,omitempty"` +} + +// ReadResourceRequest 读取资源请求 +type ReadResourceRequest struct { + URI string `json:"uri"` +} + +// ReadResourceResponse 读取资源响应 +type ReadResourceResponse struct { + Contents []ResourceContent `json:"contents"` +} + +// ResourceContent 资源内容 +type ResourceContent struct { + URI string `json:"uri"` + MimeType string `json:"mimeType,omitempty"` + Text string `json:"text,omitempty"` + Blob string `json:"blob,omitempty"` +} + +// SamplingRequest 采样请求 +type SamplingRequest struct { + Messages []SamplingMessage `json:"messages"` + Model string `json:"model,omitempty"` + MaxTokens int `json:"maxTokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` +} + +// SamplingMessage 采样消息 +type SamplingMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// SamplingResponse 采样响应 +type SamplingResponse struct { + Content []SamplingContent `json:"content"` + Model string `json:"model,omitempty"` + StopReason string `json:"stopReason,omitempty"` +} + +// SamplingContent 采样内容 +type SamplingContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// ToolResultPlainText 拼接工具结果中的文本（手动终止时作为「工具原始输出」）。 +func ToolResultPlainText(r *ToolResult) string { + if r == nil || len(r.Content) == 0 { + return "" + } + var b strings.Builder + for _, c := range r.Content { + b.WriteString(c.Text) + } + return strings.TrimSpace(b.String()) +} + +// AbortNoteBannerForModel 标出后续文本来自「用户手动终止工具时在弹窗中填写」，避免与 stdout/stderr 混淆。 +const AbortNoteBannerForModel = "---\n" + + "【用户终止说明｜USER INTERRUPT NOTE】\n" + + "（以下由操作者填写，用于指示模型如何继续；不是工具原始输出。）\n" + + "（Written by the operator when stopping this tool; not raw tool output.）\n" + + "---" + +// MergePartialToolOutputAndAbortNote 格式：工具原始输出 + 醒目标题 + 用户终止说明（无说明则原样返回 partial）。 +func MergePartialToolOutputAndAbortNote(partial, userNote string) string { + partial = strings.TrimSpace(partial) + userNote = strings.TrimSpace(userNote) + if userNote == "" { + return partial + } + section := AbortNoteBannerForModel + "\n" + userNote + if partial == "" { + return section + } + return partial + "\n\n" + section +} diff --git a/internal/projectprompt/blackboard.go b/internal/projectprompt/blackboard.go new file mode 100644 index 00000000..d3e3ae76 --- /dev/null +++ b/internal/projectprompt/blackboard.go @@ -0,0 +1,132 @@ +// Package projectprompt 提供项目黑板相关的系统提示文本（纯字符串，无 database 依赖）。 +// 供 agent / multiagent 等包引用，避免 agent → project 导入环导致 gopls 元数据失败。 +package projectprompt + +import ( + "strings" + + "cyberstrike-ai/internal/mcp/builtin" +) + +const ( + factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知（开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化）后，**立即**调用 `upsert_project_fact`（同 fact_key 覆盖更新）。每**验证**出一条可复现漏洞（含 POC/影响）后，**立即**调用 `record_vulnerability`；与事实可各记一次。继续下一步工作前优先落库，避免上下文压缩后细节丢失。未绑项目时说明无法写黑板，仍在本轮保留证据摘要。" + factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时，由协调者及时写入，勿假定子代理已记。" + factRhythmSubAgentSuffix = "若工具集中无上述工具，须在交付物末尾给出「待落库」结构化条目（fact_key 建议、summary、body/POC 要点），供协调者**立即**写入。" +) + +// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏（Markdown，供 agents/*.md 与文档对齐）。 +func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string { + var b strings.Builder + b.WriteString("- **边渗透边记录（强制节奏）**：") + b.WriteString(factRhythmCore) + if coordinator { + b.WriteString(factRhythmCoordinatorSuffix) + } + if subAgent { + b.WriteString(factRhythmSubAgentSuffix) + } + return b.String() +} + +func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string { + var b strings.Builder + b.WriteString("- **边渗透边记录（强制节奏）**：勿等会话结束或收尾再批量写入。每**确认**一条新认知（开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化）后，**立即**调用 ") + b.WriteString(builtin.ToolUpsertProjectFact) + b.WriteString("（同 fact_key 覆盖更新）。每**验证**出一条可复现漏洞（含 POC/影响）后，**立即**调用 ") + b.WriteString(builtin.ToolRecordVulnerability) + b.WriteString("；与事实可各记一次。继续下一步工作前优先落库，避免上下文压缩后细节丢失。未绑项目时说明无法写黑板，仍在本轮保留证据摘要。") + if coordinator { + b.WriteString(factRhythmCoordinatorSuffix) + } + if subAgent { + b.WriteString(factRhythmSubAgentSuffix) + } + return b.String() +} + +func factEdgeRecordingGuidance() string { + return `### 事实关系边（links） + +- 写入 **finding / chain / exploit / poc** 时，**必须**在 ` + "`upsert_project_fact`" + ` 中提供 ` + "`links`" + `（**推荐 ` + "`from`" + `**：来源 fact 指向当前 fact，即 ` + "`from`" + ` → 当前 ` + "`fact_key`" + `）。 +- **最少要求**：finding 类至少 1 条 from=target/* + type=discovered_on（即 target → finding）；在 finding 上记录 exploit 用 from=exploit/* + type=exploits（即 exploit → finding）。 +- **常用 type**：` + "`discovered_on`" + `（发现在哪）、` + "`depends_on`" + `（复现前置）、` + "`leads_to`" + `（认知推进）、` + "`enables`" + `（扩大攻击面）、` + "`exploits`" + `（利用关系）、` + "`contains`" + `（资产包含）、` + "`part_of`" + `（属于链/组）、` + "`supports`" + `（证据支撑）。 +- 更新时：**省略 links 保留已有边**；传入 links 则**替换**全部关系边（from → 当前 fact）。 +- body 中「依赖事实」段落可与 links 并存（人读）；结构化关系以 links 为准。` +} + +func factRecordingGuidanceBlock() string { + return `### 事实写入规范（审计复现 / 知识沉淀） + +- **summary**：索引用一行，须含「什么 + 在哪 + 如何触发/验证」要点，禁止只写结论（如仅写「存在 SQLi」）。 +- **body**：完整可复现上下文，写入 ` + "`upsert_project_fact`" + ` 的 body 字段；索引不含 body，后续会话须靠 ` + "`get_project_fact`" + ` 取回。 +- **category / fact_key 建议**： + - 环境认知：` + "`target/`" + `、` + "`auth/`" + `、` + "`infra/`" + `、` + "`business/`" + `（body 用环境模板即可） + - 发现与利用：` + "`finding/`" + `、` + "`chain/`" + `、` + "`exploit/`" + `、` + "`poc/`" + `（**必须**用攻击链模板填满 body：入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID） +- **与漏洞记录分工**：` + "`record_vulnerability`" + ` 记可交付 findings；事实记**复现所需的全部上下文**（含失败尝试、绕过、依赖会话），二者可各记一次。 +- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入，勿散落多个 key 导致上下文丢失。` +} + +// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块（单/多 Agent 主代理共用）。 +func FactRecordingBlackboardSection(coordinatorDelegate bool) string { + var b strings.Builder + b.WriteString("## 项目黑板（事实）与漏洞记录（分离）\n\n") + b.WriteString("当前对话若已绑定项目，系统会自动注入「项目黑板索引」（仅 fact_key + 摘要）。**摘要不足时必须调用 ") + b.WriteString(builtin.ToolGetProjectFact) + b.WriteString("(fact_key) 获取 body，禁止凭摘要臆造细节。**\n\n") + b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false)) + b.WriteString("\n\n") + b.WriteString("- **环境/目标/认证等认知**（非正式漏洞条目）：使用 ") + b.WriteString(builtin.ToolUpsertProjectFact) + b.WriteString("，fact_key 建议 `category/slug`（如 target/primary_domain），同 key 覆盖更新；body 记端口/版本/凭据特征与证据来源。\n") + b.WriteString("- **发现与利用上下文**（审计复现）：fact_key 建议 finding/、chain/、exploit/、poc/ 前缀；**body 必填**完整攻击链（入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id），**禁止仅写结论**；summary 写「什么 + 在哪 + 如何验证」一行要点。\n") + b.WriteString("- **可交付漏洞**：使用 ") + b.WriteString(builtin.ToolRecordVulnerability) + b.WriteString("，含标题、严重程度、类型、目标、证明（POC）、影响、修复建议。记前可先 ") + b.WriteString(builtin.ToolListVulnerabilities) + b.WriteString(" 查重，详情用 ") + b.WriteString(builtin.ToolGetVulnerability) + b.WriteString("(id)（默认仅当前项目/会话）。\n") + b.WriteString("- 同一发现可能需**各记一次**（事实记**完整攻击链与 exploit 细节**供复现，漏洞记正式 findings）。误报用 ") + b.WriteString(builtin.ToolDeprecateProjectFact) + b.WriteString(" 或漏洞状态 false_positive。\n") + b.WriteString("- 事实多时用 ") + b.WriteString(builtin.ToolListProjectFacts) + b.WriteString(" / ") + b.WriteString(builtin.ToolSearchProjectFacts) + b.WriteString(" 检索。\n\n") + b.WriteString(factEdgeRecordingGuidance()) + b.WriteString("\n\n") + b.WriteString(factRecordingGuidanceBlock()) + b.WriteString("\n\n严重程度：critical / high / medium / low / info。证明须含足够证据（请求响应、截图、命令输出等）。") + return b.String() +} + +// FactRecordingSubAgentSection 子代理边渗透边记录（无工具时输出待落库条目）。 +func FactRecordingSubAgentSection() string { + return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n" +} + +// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown（工具名为字面量，供 agents/*.md）。 +func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string { + var b strings.Builder + b.WriteString("## 项目黑板（事实）与漏洞记录（分离）\n\n") + b.WriteString("当前对话若已绑定项目，系统会自动注入「项目黑板索引」（仅 `fact_key` + 摘要）。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body，禁止凭摘要臆造细节。**\n\n") + b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false)) + b.WriteString("\n\n") + b.WriteString("- **环境/目标/认证等认知**（非正式漏洞）：使用 **`upsert_project_fact`**，`fact_key` 建议 `category/slug`（如 `target/primary_domain`），同 key 覆盖更新；body 记端口/版本/凭据特征与证据来源。\n") + b.WriteString("- **发现与利用上下文**（审计复现）：`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀；**body 必填**完整攻击链（入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`），**禁止仅写结论**；summary 写「什么 + 在哪 + 如何验证」一行要点。\n") + b.WriteString("- **可交付漏洞**：使用 **`record_vulnerability`**（标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议）。严重程度 critical / high / medium / low / info。\n") + b.WriteString("- 同一发现可能需**各记一次**（事实记可复现攻击链，漏洞记正式 findings）。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n") + b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n") + b.WriteString(factEdgeRecordingGuidance()) + b.WriteString("\n\n") + b.WriteString(factRecordingGuidanceBlock()) + b.WriteString("\n\n严重程度：critical / high / medium / low / info。证明须含足够证据（请求响应、截图、命令输出等）。") + return b.String() +} + +// FactEdgeRecordingGuidance 写入边时的 Agent 规范（供 project 包复用）。 +func FactEdgeRecordingGuidance() string { return factEdgeRecordingGuidance() } + +// FactRecordingGuidanceBlock 事实写入规范块（供 project 包复用）。 +func FactRecordingGuidanceBlock() string { return factRecordingGuidanceBlock() } diff --git a/internal/projectprompt/shell_tools.go b/internal/projectprompt/shell_tools.go new file mode 100644 index 00000000..6314a181 --- /dev/null +++ b/internal/projectprompt/shell_tools.go @@ -0,0 +1,11 @@ +package projectprompt + +// ShellExecExecuteGuidanceSection 供单代理/多代理系统提示追加：exec 与 execute 分工（尽量短）。 +func ShellExecExecuteGuidanceSection() string { + return `Shell（exec/execute）：有专用 MCP 工具时优先专用工具；系统命令（管道、workdir、后台 &）用 exec；skills/ 内脚本（配合 read_file、skill）用 execute；多步扫描分拆调用，禁止一条 shell 串多个扫描器。长脚本、请求体或 Payload 必须先用 write_file 写入会话工作目录，再用 exec/execute 执行短命令；禁止把长内容嵌入 command。下载/临时文件须写入系统提示中的「会话工作目录」，禁止用 /tmp。` +} + +// ShellExecExecuteGuidanceReconSuffix 侦察子代理可选追加（一行）。 +func ShellExecExecuteGuidanceReconSuffix() string { + return `枚举优先 subfinder、amass 等专用 MCP，勿 exec/execute 拼长链。` +} diff --git a/internal/reasoning/eino.go b/internal/reasoning/eino.go new file mode 100644 index 00000000..3d7c19ae --- /dev/null +++ b/internal/reasoning/eino.go @@ -0,0 +1,428 @@ +// Package reasoning maps user/config intent to CloudWeGo Eino OpenAI ChatModel fields +// (ReasoningEffort, ExtraFields such as thinking / reasoning_effort / output_config). +package reasoning + +import ( + "strings" + + "cyberstrike-ai/internal/config" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" +) + +// ClientIntent is optional per-request override from ChatRequest.reasoning. +type ClientIntent struct { + Mode string + Effort string +} + +type wireProfile int + +const ( + wireNone wireProfile = iota + wireClaude + wireDeepseek + wireOpenAI + wireOutputConfig +) + +// ApplyPlanExecutePlannerModelConfig configures the plan_execute planner/replanner +// ChatModel. Those Eino agents call WithToolChoice(Forced); several gateways reject +// thinking / reasoning fields on the same request (tool_choice required/object). +// Executor should keep the normal ApplyToEinoChatModelConfig path. +func ApplyPlanExecutePlannerModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig) { + if cfg == nil || oa == nil { + return + } + mergeExtraRequestFields(cfg, oa.Reasoning.ExtraRequestFields) + clearReasoningFromChatModelConfig(cfg) + if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + // DeepSeek enables thinking by default, so omission would not actually + // disable it for the planner's forced tool-choice requests. + applyThinkingDisabled(cfg) + } +} + +func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) { + if cfg == nil { + return + } + cfg.ReasoningEffort = "" + if cfg.ExtraFields != nil { + for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} { + delete(cfg.ExtraFields, key) + } + if len(cfg.ExtraFields) == 0 { + cfg.ExtraFields = nil + } + } +} + +func mergeExtraRequestFields(cfg *einoopenai.ChatModelConfig, fields map[string]interface{}) { + if cfg == nil || len(fields) == 0 { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any, len(fields)) + } + for k, v := range fields { + cfg.ExtraFields[k] = v + } +} + +// ApplyToEinoChatModelConfig merges reasoning-related options into cfg. +// Precondition: cfg already has APIKey, BaseURL, Model, HTTPClient set. +func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig, client *ClientIntent) { + if cfg == nil || oa == nil { + return + } + sr := &oa.Reasoning + allowClient := sr.AllowClientReasoningEffective() + mode := effectiveMode(sr, client, allowClient) + + // Admin-defined root fields are independent of the selected reasoning wire + // profile. Merge them first so mode=off can remove only reasoning controls + // while preserving unrelated gateway options. + mergeExtraRequestFields(cfg, sr.ExtraRequestFields) + if mode == "off" { + clearReasoningFromChatModelConfig(cfg) + // Strict OpenAI endpoints reject unknown `thinking` fields, whereas the + // DeepSeek API enables thinking by default and requires an explicit + // thinking.type=disabled switch. Detect the actual DeepSeek target even + // when the configured reasoning profile was left as openai_compat. + if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + applyThinkingDisabled(cfg) + } + return + } + + // Claude (Anthropic): merge admin extras first; optional extended thinking maps to top-level `thinking`. + // DeepSeek/OpenAI-style fields are not sent. + if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") || + strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") { + applyClaudeExtendedThinking(cfg, mode, effectiveEffort(sr, client, allowClient), oa.Model) + return + } + + effort := effectiveEffort(sr, client, allowClient) + prof := resolveWireProfile(oa, sr) + + switch prof { + case wireClaude, wireNone: + return + case wireDeepseek: + applyDeepseek(cfg, mode, effort) + case wireOutputConfig: + applyOutputConfigEffort(cfg, mode, effort) + default: // wireOpenAI + applyOpenAICompat(cfg, mode, effort) + } +} + +// AgenticOpenAIExtraFields returns reasoning-related request fields for +// agenticopenai.ChatConfig. The agentic chat backend currently exposes provider +// extensions through ExtraFields instead of typed ReasoningEffort fields. +func AgenticOpenAIExtraFields(oa *config.OpenAIConfig, client *ClientIntent) map[string]any { + if oa == nil { + return nil + } + sr := &oa.Reasoning + allowClient := sr.AllowClientReasoningEffective() + mode := effectiveMode(sr, client, allowClient) + fields := cloneExtraRequestFields(sr.ExtraRequestFields) + if mode == "off" { + clearReasoningExtraFields(fields) + if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "disabled"} + } + return fields + } + if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") || + strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") { + return fields + } + effort := effectiveEffort(sr, client, allowClient) + switch resolveWireProfile(oa, sr) { + case wireDeepseek: + if mode == "auto" || mode == "on" { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "enabled"} + } + if effort != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["reasoning_effort"] = effortStringForAPI(effort) + } + case wireOutputConfig: + e := effort + if mode == "on" && e == "" { + e = "high" + } + if e != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["output_config"] = map[string]any{"effort": effortStringForAPI(e)} + } + default: + e := effort + if mode == "on" && e == "" { + e = "medium" + } + if e != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["reasoning_effort"] = effortStringForAPI(e) + } + } + return fields +} + +// AgenticOpenAIPlannerExtraFields mirrors ApplyPlanExecutePlannerModelConfig for +// agenticopenai.ChatConfig: keep admin extras, strip reasoning controls, and +// explicitly disable DeepSeek thinking where omission would still think. +func AgenticOpenAIPlannerExtraFields(oa *config.OpenAIConfig) map[string]any { + if oa == nil { + return nil + } + fields := cloneExtraRequestFields(oa.Reasoning.ExtraRequestFields) + clearReasoningExtraFields(fields) + if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "disabled"} + } + return fields +} + +func cloneExtraRequestFields(fields map[string]interface{}) map[string]any { + if len(fields) == 0 { + return nil + } + out := make(map[string]any, len(fields)) + for k, v := range fields { + out[k] = v + } + return out +} + +func clearReasoningExtraFields(fields map[string]any) { + for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} { + delete(fields, key) + } +} + +// applyClaudeExtendedThinking sets Anthropic Messages API fields per official guidance: +// - Adaptive models (4.6+): thinking.type=adaptive; output_config.effort only when user sets effort (API default is high). +// - Sonnet 3.7: thinking.type=enabled + budget_tokens=10000 (doc example); effort is not mapped — use extra_request_fields for custom budget. +func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) { + if cfg == nil || mode == "off" { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + m := strings.ToLower(strings.TrimSpace(model)) + sonnet37 := isClaudeSonnet37(m) + + if _, exists := cfg.ExtraFields["thinking"]; !exists { + cfg.ExtraFields["thinking"] = claudeThinkingForModel(m, sonnet37) + } + + applyClaudeOutputConfigEffort(cfg, effort, sonnet37) +} + +// claudeSonnet37DefaultBudgetTokens matches Anthropic extended-thinking documentation examples (budget_tokens with max_tokens 16000). +const claudeSonnet37DefaultBudgetTokens = 10000 + +func isClaudeSonnet37(m string) bool { + return strings.Contains(m, "claude-3-7-sonnet") || + strings.Contains(m, "3-7-sonnet") || + strings.Contains(m, "sonnet-3.7") +} + +func claudeThinkingForModel(m string, sonnet37 bool) map[string]any { + if sonnet37 { + return map[string]any{ + "type": "enabled", + "budget_tokens": claudeSonnet37DefaultBudgetTokens, + "display": "summarized", + } + } + // Opus 4.7+: manual enabled+budget rejected — adaptive only. + if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") { + return map[string]any{ + "type": "adaptive", + "display": "summarized", + } + } + return map[string]any{ + "type": "adaptive", + "display": "summarized", + } +} + +// applyClaudeOutputConfigEffort sets top-level output_config.effort only when effort is explicitly configured. +// Omitted effort uses the API default (high); do not inject effort on mode:on alone. +func applyClaudeOutputConfigEffort(cfg *einoopenai.ChatModelConfig, effort string, sonnet37 bool) { + if cfg == nil || sonnet37 { + return + } + if _, exists := cfg.ExtraFields["output_config"]; exists { + return + } + e := effortStringForAPI(effort) + if e == "" { + return + } + cfg.ExtraFields["output_config"] = map[string]any{"effort": e} +} + +func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string { + server := strings.ToLower(strings.TrimSpace(sr.ModeEffective())) + if server == "" || server == "default" { + server = "auto" + } + if !allowClient || client == nil { + return server + } + cm := strings.ToLower(strings.TrimSpace(client.Mode)) + if cm == "" || cm == "default" { + return server + } + return cm +} + +func effectiveEffort(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string { + se := normalizeEffort(sr.Effort) + if !allowClient || client == nil { + return se + } + ce := normalizeEffort(client.Effort) + if ce != "" { + return ce + } + return se +} + +func normalizeEffort(s string) string { + e := strings.ToLower(strings.TrimSpace(s)) + switch e { + case "low", "medium", "high", "max", "xhigh": + return e + default: + return "" + } +} + +// usesExtraFieldsReasoningEffort 为 Eino 无枚举的最高档 effort，经 ExtraFields 原样下发（max / xhigh 由网关自行识别，不做互转）。 +func usesExtraFieldsReasoningEffort(e string) bool { + return e == "max" || e == "xhigh" +} + +func resolveWireProfile(oa *config.OpenAIConfig, sr *config.OpenAIReasoningConfig) wireProfile { + provider := strings.TrimSpace(oa.Provider) + if strings.EqualFold(provider, "claude") || strings.EqualFold(provider, "anthropic") { + return wireClaude + } + p := strings.ToLower(strings.TrimSpace(sr.ProfileEffective())) + switch p { + case "output_config", "output_config_effort": + return wireOutputConfig + case "openai", "openai_compat": + return wireOpenAI + case "deepseek", "deepseek_compat": + return wireDeepseek + case "auto", "": + if oa.IsDeepSeekEndpointOrModel() { + return wireDeepseek + } + return wireOpenAI + default: + return wireOpenAI + } +} + +func applyThinkingDisabled(cfg *einoopenai.ChatModelConfig) { + if cfg == nil { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["thinking"] = map[string]any{"type": "disabled"} +} + +func applyDeepseek(cfg *einoopenai.ChatModelConfig, mode, effort string) { + // auto: enable thinking for DeepSeek line; on: same; auto without effort still opens thinking. + if mode == "auto" || mode == "on" { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["thinking"] = map[string]any{"type": "enabled"} + } + if effort != "" { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(effort) + } +} + +func applyOpenAICompat(cfg *einoopenai.ChatModelConfig, mode, effort string) { + if mode == "auto" && effort == "" { + return + } + e := effort + if mode == "on" && e == "" { + e = "medium" + } + if e == "" { + return + } + if usesExtraFieldsReasoningEffort(e) { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(e) + return + } + switch e { + case "low": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelLow + case "medium": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelMedium + case "high": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelHigh + } +} + +func applyOutputConfigEffort(cfg *einoopenai.ChatModelConfig, mode, effort string) { + if mode == "auto" && effort == "" { + return + } + e := effort + if mode == "on" && e == "" { + e = "high" + } + if e == "" { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["output_config"] = map[string]any{"effort": effortStringForAPI(e)} +} + +func effortStringForAPI(e string) string { + // 原样透传：OpenAI 官方多为 xhigh，部分兼容网关为 max，由配置/对话 effort 选择。 + return strings.ToLower(strings.TrimSpace(e)) +} diff --git a/internal/reasoning/eino_test.go b/internal/reasoning/eino_test.go new file mode 100644 index 00000000..c9eadf8a --- /dev/null +++ b/internal/reasoning/eino_test.go @@ -0,0 +1,424 @@ +package reasoning + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/config" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/schema" +) + +var reasoningPayloadKeysForTest = []string{"thinking", "reasoning_effort", "output_config", "reasoning"} + +func assertNoReasoningFields(t *testing.T, cfg *einoopenai.ChatModelConfig) { + t.Helper() + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + for _, key := range reasoningPayloadKeysForTest { + if _, ok := cfg.ExtraFields[key]; ok { + t.Fatalf("expected %q omitted, got %#v", key, cfg.ExtraFields) + } + } +} + +func TestEffortStringForAPI_passthrough(t *testing.T) { + cases := map[string]string{ + "max": "max", + "xhigh": "xhigh", + "HIGH": "high", + "Medium": "medium", + } + for in, want := range cases { + if got := effortStringForAPI(in); got != want { + t.Fatalf("%q -> %q, want %q", in, got, want) + } + } +} + +func TestNormalizeEffort_maxAndXhigh(t *testing.T) { + if normalizeEffort("xhigh") != "xhigh" { + t.Fatal("xhigh not accepted") + } + if normalizeEffort("max") != "max" { + t.Fatal("max not accepted") + } +} + +func TestApplyOpenAICompat_xhighExtraField(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "xhigh", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if cfg.ExtraFields == nil { + t.Fatal("expected ExtraFields") + } + if got, _ := cfg.ExtraFields["reasoning_effort"].(string); got != "xhigh" { + t.Fatalf("reasoning_effort=%q", got) + } +} + +func TestAgenticOpenAIExtraFields_openAICompatReasoningEffort(t *testing.T) { + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIExtraFields(oa, nil) + if got["reasoning_effort"] != "high" { + t.Fatalf("reasoning_effort=%#v, want high in %#v", got["reasoning_effort"], got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIExtraFields_reasoningOffPreservesUnrelatedFields(t *testing.T) { + oa := &config.OpenAIConfig{ + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "off", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "thinking": map[string]any{"type": "enabled"}, + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIExtraFields(oa, nil) + for _, key := range reasoningPayloadKeysForTest { + if _, ok := got[key]; ok { + t.Fatalf("agentic fields unexpectedly contain %q: %#v", key, got) + } + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIPlannerExtraFields_deepseekDisablesThinking(t *testing.T) { + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com", + Model: "deepseek-chat", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "auto", + Mode: "on", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIPlannerExtraFields(oa) + if got["reasoning_effort"] != nil { + t.Fatalf("planner should strip reasoning_effort: %#v", got) + } + thinking, ok := got["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled, got %#v", got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIPlannerExtraFields_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) { + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4-flash", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIPlannerExtraFields(oa) + if _, ok := got["reasoning_effort"]; ok { + t.Fatalf("planner should strip reasoning_effort: %#v", got) + } + thinking, ok := got["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://antchat.example.com/v1", + Model: "minimax-m3", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + }, + } + ApplyPlanExecutePlannerModelConfig(cfg, oa) + assertNoReasoningFields(t, cfg) + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyPlanExecutePlannerModelConfig_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4-flash", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + }, + } + ApplyPlanExecutePlannerModelConfig(cfg, oa) + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + if _, ok := cfg.ExtraFields["reasoning_effort"]; ok { + t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields) + } + thinking, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", cfg.ExtraFields) + } + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyReasoningOff_omitsAllReasoningFields(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "output_config": map[string]any{"effort": "high"}, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", + Effort: "high", + Profile: "openai_compat", + ExtraRequestFields: map[string]interface{}{ + "thinking": map[string]any{"type": "disabled"}, + "reasoning": map[string]any{"effort": "high"}, + "vendor_option": true, + }, + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + assertNoReasoningFields(t, cfg) + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyReasoningOff_clientOverrideOmit(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", Effort: "high", Profile: "openai_compat", + }} + ApplyToEinoChatModelConfig(cfg, oa, &ClientIntent{Mode: "off", Effort: "high"}) + assertNoReasoningFields(t, cfg) +} + +func TestApplyReasoningOff_deepseekExplicitlyDisablesDefaultThinking(t *testing.T) { + for _, profile := range []string{"deepseek_compat", "auto", "openai_compat"} { + t.Run(profile, func(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com", + Model: "deepseek-v4-pro", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", Effort: "high", Profile: profile, + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + if _, ok := cfg.ExtraFields["reasoning_effort"]; ok { + t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields) + } + thinking, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected DeepSeek thinking disabled, got %#v", cfg.ExtraFields) + } + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } + }) + } +} + +func TestApplyReasoningOff_wirePayloadOmitsThinking(t *testing.T) { + var requestBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + if err := json.Unmarshal(body, &requestBody); err != nil { + t.Errorf("decode request body: %v; body=%s", err, body) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + })) + defer srv.Close() + + cfg := &einoopenai.ChatModelConfig{ + APIKey: "test-key", + BaseURL: srv.URL, + Model: "gpt-4o-mini", + } + oa := &config.OpenAIConfig{ + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", Effort: "high", Profile: "openai_compat", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + model, err := einoopenai.NewChatModel(context.Background(), cfg) + if err != nil { + t.Fatalf("new chat model: %v", err) + } + if _, err := model.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}); err != nil { + t.Fatalf("generate: %v", err) + } + for _, key := range reasoningPayloadKeysForTest { + if _, ok := requestBody[key]; ok { + t.Fatalf("wire payload unexpectedly contains %q: %#v", key, requestBody) + } + } +} + +func TestApplyOpenAICompat_maxPassthrough(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "max", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + got, _ := cfg.ExtraFields["reasoning_effort"].(string) + if got != "max" { + t.Fatalf("max effort wire=%q, want max", got) + } +} + +func TestApplyClaude_adaptiveOutputConfigEffort(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-opus-4-8", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + Effort: "xhigh", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + th, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || th["type"] != "adaptive" { + t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"]) + } + oc, ok := cfg.ExtraFields["output_config"].(map[string]any) + if !ok { + t.Fatal("expected output_config") + } + if oc["effort"] != "xhigh" { + t.Fatalf("effort=%v", oc["effort"]) + } +} + +func TestApplyClaude_sonnet37OfficialBudget(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-3-7-sonnet-latest", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + Effort: "low", // 3.7 has no output_config.effort; effort is not mapped to budget_tokens + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + th, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || th["type"] != "enabled" { + t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"]) + } + if th["budget_tokens"] != claudeSonnet37DefaultBudgetTokens { + t.Fatalf("budget_tokens=%v, want official example %d", th["budget_tokens"], claudeSonnet37DefaultBudgetTokens) + } + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("sonnet 3.7 should not set output_config") + } +} + +func TestApplyClaude_onWithoutEffortOmitsOutputConfig(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-sonnet-4-6", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("on without explicit effort should omit output_config (API default high)") + } +} + +func TestApplyClaude_autoWithoutEffortSkipsOutputConfig(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-sonnet-4-6", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "auto", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("auto without effort should omit output_config") + } +} diff --git a/internal/tooloutput/spill.go b/internal/tooloutput/spill.go new file mode 100644 index 00000000..1352e219 --- /dev/null +++ b/internal/tooloutput/spill.go @@ -0,0 +1,292 @@ +// Package tooloutput spills oversized tool stdout/results to local files under +// the reduction cache tree (tmp/reduction/...), so agents can read_file the +// full text after context truncation. +package tooloutput + +import ( + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "unicode/utf8" + + "github.com/google/uuid" +) + +const ( + defaultRootDir = "tmp/reduction" + readFileHint = "read_file" +) + +// SpillOpts scopes where a trunc file is written (mirrors reduction RootDir layout). +type SpillOpts struct { + RootDir string // reduction_root_dir or empty → tmp/reduction + ProjectID string + ConversationID string + ExecutionID string // preferred file name; empty → uuid +} + +// SessionRoot returns the conversation/project-scoped reduction cache root. +func SessionRoot(configuredBase, projectID, conversationID string) string { + base := strings.TrimSpace(configuredBase) + if base == "" { + base = defaultRootDir + } + if pid := strings.TrimSpace(projectID); pid != "" { + return filepath.Join(base, "projects", sanitizeSegment(pid)) + } + conv := strings.TrimSpace(conversationID) + if conv == "" { + conv = "default" + } + return filepath.Join(base, "conversations", sanitizeSegment(conv)) +} + +// WriteTruncFile writes full content under {sessionRoot}/trunc/{id} and returns +// an absolute path suitable for read_file. +func WriteTruncFile(opts SpillOpts, content string) (string, error) { + session := SessionRoot(opts.RootDir, opts.ProjectID, opts.ConversationID) + id := strings.TrimSpace(opts.ExecutionID) + if id == "" { + id = uuid.NewString() + } + id = sanitizeSegment(id) + dir := filepath.Join(session, "trunc") + if err := os.MkdirAll(dir, 0o755); err != nil { + return "", fmt.Errorf("mkdir tool output trunc dir: %w", err) + } + path := filepath.Join(dir, id) + if abs, err := filepath.Abs(path); err == nil { + path = abs + } + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + return "", fmt.Errorf("write tool output trunc file: %w", err) + } + return path, nil +} + +// BoundWithSpill truncates full text into a notice after +// spilling the original to disk. The returned string is always ≤ maxBytes when +// maxBytes > 0. On spill failure it falls back to a prefix + marker (no path). +func BoundWithSpill(full string, maxBytes int, opts SpillOpts) string { + if maxBytes <= 0 || len(full) <= maxBytes { + return full + } + path, err := WriteTruncFile(opts, full) + if err != nil { + return boundPrefixOnly(full, maxBytes, len(full), "") + } + return FormatPersistedOutput(full, path, maxBytes) +} + +// FormatPersistedOutput builds a reduction-compatible notice with head/tail +// previews that fits in maxBytes. +func FormatPersistedOutput(full, filePath string, maxBytes int) string { + return formatPersisted(len(full), filePath, full, maxBytes) +} + +// FormatPersistedFromFile builds the notice using previews read from an already +// spilled file (streaming collectors that never kept the full string in memory). +func FormatPersistedFromFile(filePath string, originalSize, maxBytes int) string { + previewSrc := "" + if data, err := os.ReadFile(filePath); err == nil { + previewSrc = string(data) + if originalSize <= 0 { + originalSize = len(data) + } + } + return formatPersisted(originalSize, filePath, previewSrc, maxBytes) +} + +func formatPersisted(originalSize int, filePath, previewSrc string, maxBytes int) string { + if maxBytes <= 0 { + maxBytes = 12000 + } + // Always keep the absolute path readable for read_file, even under tight budgets. + minimal := fmt.Sprintf( + "\nOutput too large (%d). Full output saved to: %s\nUse %s to read.\n", + originalSize, filePath, readFileHint, + ) + if len(minimal) > maxBytes { + core := fmt.Sprintf("Full output saved to: %s", filePath) + if len(core) <= maxBytes { + return core + } + // Path longer than budget: keep as much of the path as possible after a short prefix. + prefix := "Full output saved to: " + suffix := "" + room := maxBytes - len(prefix) - len(suffix) + if room <= 0 { + return clampPrefix(core, maxBytes) + } + return prefix + clampSuffix(filePath, room) + suffix + } + + previewBudget := maxBytes - len(minimal) + 32 // approximate room beyond minimal shell + if previewBudget > 4000 { + previewBudget = 4000 + } + if previewBudget < 0 { + previewBudget = 0 + } + for previewBudget >= 0 { + half := previewBudget / 2 + head := clampPrefix(previewSrc, half) + tail := clampSuffix(previewSrc, previewBudget-half) + notice := fmt.Sprintf( + "\nOutput too large (%d). Full output saved to: %s\nUse %s with offset/limit to read parts of the file.\nPreview (first %d):\n%s\n\nPreview (last %d):\n%s\n\n", + originalSize, filePath, readFileHint, len(head), head, len(tail), tail, + ) + if len(notice) <= maxBytes { + return notice + } + if previewBudget == 0 { + return minimal + } + previewBudget = previewBudget * 3 / 4 + } + return minimal +} + +func boundPrefixOnly(full string, maxBytes, originalSize int, filePath string) string { + marker := fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes]...", originalSize, maxBytes) + if filePath != "" { + marker = fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes; full output: %s]...", originalSize, maxBytes, filePath) + } + budget := maxBytes - len(marker) + if budget < 0 { + return clampPrefix(marker, maxBytes) + } + return clampPrefix(full, budget) + marker +} + +func clampPrefix(s string, n int) string { + if n <= 0 { + return "" + } + if len(s) <= n { + return s + } + for n > 0 && !utf8.RuneStart(s[n]) { + n-- + } + return s[:n] +} + +func clampSuffix(s string, n int) string { + if n <= 0 { + return "" + } + if len(s) <= n { + return s + } + start := len(s) - n + for start < len(s) && !utf8.RuneStart(s[start]) { + start++ + } + return s[start:] +} + +func sanitizeSegment(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return "default" + } + s = strings.ReplaceAll(s, string(filepath.Separator), "-") + s = strings.ReplaceAll(s, "/", "-") + s = strings.ReplaceAll(s, "\\", "-") + s = strings.ReplaceAll(s, "..", "__") + if len(s) > 180 { + s = s[:180] + } + return s +} + +// Tee writes every byte to a trunc file while callers keep only a bounded +// in-memory prefix. Safe for concurrent stdout/stderr writers. +type Tee struct { + mu sync.Mutex + opts SpillOpts + file *os.File + path string + err error + open bool +} + +// NewTee prepares a lazy spill file (created on first Write). +func NewTee(opts SpillOpts) *Tee { + return &Tee{opts: opts} +} + +// Write appends to the spill file, creating it on first use. +func (t *Tee) Write(p []byte) (int, error) { + if t == nil { + return len(p), nil + } + t.mu.Lock() + defer t.mu.Unlock() + if err := t.ensureOpenLocked(); err != nil { + return len(p), nil // best-effort: never fail the tool pipe + } + if t.file == nil { + return len(p), nil + } + _, _ = t.file.Write(p) + return len(p), nil +} + +func (t *Tee) ensureOpenLocked() error { + if t.open || t.err != nil { + return t.err + } + t.open = true + session := SessionRoot(t.opts.RootDir, t.opts.ProjectID, t.opts.ConversationID) + id := strings.TrimSpace(t.opts.ExecutionID) + if id == "" { + id = uuid.NewString() + } + id = sanitizeSegment(id) + dir := filepath.Join(session, "trunc") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.err = err + return err + } + path := filepath.Join(dir, id) + if abs, err := filepath.Abs(path); err == nil { + path = abs + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600) + if err != nil { + t.err = err + return err + } + t.file = f + t.path = path + return nil +} + +// Path returns the absolute spill path after any Write (may be empty if unused/failed). +func (t *Tee) Path() string { + if t == nil { + return "" + } + t.mu.Lock() + defer t.mu.Unlock() + return t.path +} + +// Close flushes and closes the spill file. +func (t *Tee) Close() error { + if t == nil { + return nil + } + t.mu.Lock() + defer t.mu.Unlock() + if t.file == nil { + return nil + } + err := t.file.Close() + t.file = nil + return err +} diff --git a/internal/tooloutput/spill_test.go b/internal/tooloutput/spill_test.go new file mode 100644 index 00000000..0b6c20a9 --- /dev/null +++ b/internal/tooloutput/spill_test.go @@ -0,0 +1,58 @@ +package tooloutput + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestBoundWithSpillWritesFullFile(t *testing.T) { + root := t.TempDir() + full := strings.Repeat("A", 2000) + "TAIL" + out := BoundWithSpill(full, 512, SpillOpts{ + RootDir: root, + ConversationID: "conv-1", + ExecutionID: "exec-1", + }) + if len(out) > 512 { + t.Fatalf("bounded output exceeds max: %d", len(out)) + } + if !strings.Contains(out, "") { + t.Fatalf("expected persisted-output notice: %q", out) + } + path := filepath.Join(root, "conversations", "conv-1", "trunc", "exec-1") + abs, err := filepath.Abs(path) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(out, abs) { + t.Fatalf("expected absolute path %q in notice: %q", abs, out) + } + got, err := os.ReadFile(abs) + if err != nil { + t.Fatal(err) + } + if string(got) != full { + t.Fatalf("spilled content mismatch: got %d want %d", len(got), len(full)) + } +} + +func TestTeeThenFormatPersistedFromFile(t *testing.T) { + root := t.TempDir() + tee := NewTee(SpillOpts{RootDir: root, ConversationID: "c", ExecutionID: "e"}) + full := strings.Repeat("xy", 100) + if _, err := tee.Write([]byte(full)); err != nil { + t.Fatal(err) + } + if err := tee.Close(); err != nil { + t.Fatal(err) + } + notice := FormatPersistedFromFile(tee.Path(), len(full), 512) + if len(notice) > 512 { + t.Fatalf("notice too long: %d", len(notice)) + } + if !strings.Contains(notice, tee.Path()) { + t.Fatalf("missing path in notice: %q", notice) + } +} diff --git a/internal/vision/client.go b/internal/vision/client.go new file mode 100644 index 00000000..cbfd89ee --- /dev/null +++ b/internal/vision/client.go @@ -0,0 +1,166 @@ +package vision + +import ( + "context" + "encoding/base64" + "fmt" + "net" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/llm" + "cyberstrike-ai/internal/openai" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/schema" +) + +// Client 调用独立 Vision ChatModel（单次 Generate）。 +type Client struct { + cfg config.VisionConfig + mainOA config.OpenAIConfig +} + +// NewClient 构造视觉客户端。 +func NewClient(visionCfg config.VisionConfig, mainOpenAI config.OpenAIConfig) *Client { + return &Client{cfg: visionCfg, mainOA: mainOpenAI} +} + +// Analyze 将图片字节送入 VL 模型并返回文本描述。 +func (c *Client) Analyze(ctx context.Context, img ImagePayload, question string) (string, error) { + if len(img.Bytes) == 0 { + return "", fmt.Errorf("empty image payload") + } + mime := strings.TrimSpace(img.MIMEType) + if mime == "" { + mime = "image/jpeg" + } + oa := c.cfg.OpenAICfgEffective(c.mainOA) + if strings.TrimSpace(oa.APIKey) == "" { + return "", fmt.Errorf("vision API key is empty (set vision.api_key or openai.api_key)") + } + if strings.TrimSpace(oa.Model) == "" { + return "", fmt.Errorf("vision model is empty") + } + + timeout := time.Duration(c.cfg.TimeoutSecondsEffective()) * time.Second + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + httpClient := &http.Client{ + Timeout: timeout + 15*time.Second, + Transport: &http.Transport{ + DialContext: (&net.Dialer{ + Timeout: 60 * time.Second, + KeepAlive: 60 * time.Second, + }).DialContext, + ResponseHeaderTimeout: timeout + 10*time.Second, + }, + } + + b64 := base64.StdEncoding.EncodeToString(img.Bytes) + detail := schema.ImageURLDetailLow + switch c.cfg.DetailEffective() { + case "high": + detail = schema.ImageURLDetailHigh + case "auto": + detail = schema.ImageURLDetailAuto + } + + prompt := buildVisionPrompt(question) + if llm.IsClaudeProvider(oa.Provider) { + nativeModel, err := llm.NewClaudeAgenticModel( + ctx, + oa, + httpClient, + oa.MaxCompletionTokensEffective(), + nil, + ) + if err != nil { + return "", fmt.Errorf("vision native Claude model: %w", err) + } + resp, err := nativeModel.Generate(ctx, []*schema.AgenticMessage{{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.UserInputText{Text: prompt}), + schema.NewContentBlock(&schema.UserInputImage{ + Base64Data: b64, + MIMEType: mime, + Detail: detail, + }), + }, + }}) + if err != nil { + return "", fmt.Errorf("vision native Claude generate: %w", err) + } + content, _ := llm.AgenticText(resp) + if strings.TrimSpace(content) == "" { + return "", fmt.Errorf("vision model returned empty content") + } + return strings.TrimSpace(content), nil + } + + httpClient = openai.NewEinoHTTPClient(&oa, httpClient) + maxCompletionTokens := oa.MaxCompletionTokensEffective() + modelCfg := &einoopenai.ChatModelConfig{ + APIKey: oa.APIKey, + BaseURL: strings.TrimSuffix(oa.BaseURL, "/"), + Model: oa.Model, + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + } + chatModel, err := einoopenai.NewChatModel(ctx, modelCfg) + if err != nil { + return "", fmt.Errorf("vision chat model: %w", err) + } + userMsg := &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: prompt}, + { + Type: schema.ChatMessagePartTypeImageURL, + Image: &schema.MessageInputImage{ + MessagePartCommon: schema.MessagePartCommon{ + Base64Data: &b64, + MIMEType: mime, + }, + Detail: detail, + }, + }, + }, + } + + resp, err := chatModel.Generate(ctx, []*schema.Message{userMsg}) + if err != nil { + return "", fmt.Errorf("vision generate: %w", err) + } + if resp == nil || strings.TrimSpace(resp.Content) == "" { + return "", fmt.Errorf("vision model returned empty content") + } + return strings.TrimSpace(resp.Content), nil +} + +func buildVisionPrompt(question string) string { + q := strings.TrimSpace(question) + if q == "" { + q = "请对图片做通用描述，侧重授权安全测试场景（可见文本、表单、按钮、验证码、错误信息、技术栈线索）。" + } + extra := "" + if looksLikeCaptchaQuestion(q) { + extra = "\n若为验证码：仅输出你辨认出的字符序列，不要空格、标点、解释；看不清则明确说无法识别。" + } + return `你是授权安全测试助手。请根据图片回答用户问题，只描述你能从图中确认的内容，不要编造。 +用户问题：` + q + extra +} + +func looksLikeCaptchaQuestion(q string) bool { + s := strings.ToLower(q) + for _, kw := range []string{"验证码", "captcha", "verification code", "verify code", "vcode", "图形码"} { + if strings.Contains(s, kw) { + return true + } + } + return strings.Contains(s, "只输出") && (strings.Contains(s, "字符") || strings.Contains(s, "character")) +} diff --git a/internal/vision/client_test.go b/internal/vision/client_test.go new file mode 100644 index 00000000..101aa943 --- /dev/null +++ b/internal/vision/client_test.go @@ -0,0 +1,12 @@ +package vision + +import "testing" + +func TestLooksLikeCaptchaQuestion(t *testing.T) { + if !looksLikeCaptchaQuestion("识别验证码，只输出字符") { + t.Fatal("expected captcha hint") + } + if looksLikeCaptchaQuestion("描述登录页布局") { + t.Fatal("expected non-captcha") + } +} diff --git a/internal/vision/path.go b/internal/vision/path.go new file mode 100644 index 00000000..3d9756ed --- /dev/null +++ b/internal/vision/path.go @@ -0,0 +1,72 @@ +package vision + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +var allowedImageExt = map[string]struct{}{ + ".png": {}, ".jpg": {}, ".jpeg": {}, ".webp": {}, ".gif": {}, + ".bmp": {}, ".tif": {}, ".tiff": {}, +} + +// ResolveImagePath 解析并校验可读图片路径（支持任意目录；仍校验扩展名与常规文件）。 +func ResolveImagePath(path string, cwd string) (string, error) { + p := strings.TrimSpace(path) + if p == "" { + return "", fmt.Errorf("path is empty") + } + cwdTrim := strings.TrimSpace(cwd) + if cwdTrim == "" { + var err error + cwdTrim, err = os.Getwd() + if err != nil { + return "", fmt.Errorf("getwd: %w", err) + } + } + cwdAbs, err := filepath.Abs(filepath.Clean(cwdTrim)) + if err != nil { + return "", err + } + + var candidate string + if filepath.IsAbs(p) { + candidate = filepath.Clean(p) + } else { + candidate = filepath.Clean(filepath.Join(cwdAbs, p)) + } + resolved := normalizeAbsPath(candidate) + if resolved == "" { + return "", fmt.Errorf("invalid path") + } + + ext := strings.ToLower(filepath.Ext(resolved)) + if _, ok := allowedImageExt[ext]; !ok { + return "", fmt.Errorf("unsupported image extension %q", ext) + } + + st, err := os.Stat(resolved) + if err != nil { + return "", fmt.Errorf("stat: %w", err) + } + if st.IsDir() { + return "", fmt.Errorf("not a regular file") + } + if st.Size() > 0 && st.Size() > 1<<30 { + return "", fmt.Errorf("file too large on disk") + } + return resolved, nil +} + +func normalizeAbsPath(p string) string { + abs, err := filepath.Abs(filepath.Clean(p)) + if err != nil { + return "" + } + if link, err := filepath.EvalSymlinks(abs); err == nil { + return link + } + return abs +} diff --git a/internal/vision/path_test.go b/internal/vision/path_test.go new file mode 100644 index 00000000..b38206bf --- /dev/null +++ b/internal/vision/path_test.go @@ -0,0 +1,52 @@ +package vision + +import ( + "os" + "path/filepath" + "testing" +) + +func TestResolveImagePath_underCWD(t *testing.T) { + dir := t.TempDir() + img := filepath.Join(dir, "shot.png") + if err := os.WriteFile(img, []byte{0x89, 0x50, 0x4e, 0x47}, 0o644); err != nil { + t.Fatal(err) + } + got, err := ResolveImagePath(img, dir) + if err != nil { + t.Fatal(err) + } + want := normalizeAbsPath(img) + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestResolveImagePath_absoluteOutsideCWD(t *testing.T) { + dir := t.TempDir() + cwd := t.TempDir() + img := filepath.Join(dir, "remote.png") + if err := os.WriteFile(img, []byte{0x89, 0x50, 0x4e, 0x47}, 0o644); err != nil { + t.Fatal(err) + } + got, err := ResolveImagePath(img, cwd) + if err != nil { + t.Fatalf("expected absolute path outside cwd to be allowed: %v", err) + } + want := normalizeAbsPath(img) + if got != want { + t.Fatalf("got %q want %q", got, want) + } +} + +func TestResolveImagePath_rejectsNonImageExt(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "notes.txt") + if err := os.WriteFile(f, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + _, err := ResolveImagePath(f, dir) + if err == nil { + t.Fatal("expected error for non-image extension") + } +} diff --git a/internal/vision/preprocess.go b/internal/vision/preprocess.go new file mode 100644 index 00000000..860dab63 --- /dev/null +++ b/internal/vision/preprocess.go @@ -0,0 +1,212 @@ +package vision + +import ( + "bytes" + "fmt" + "image" + "os" + "strings" + + "github.com/disintegration/imaging" +) + +// ImagePayload 送入 VL API 的图片字节与 MIME。 +type ImagePayload struct { + Bytes []byte + MIMEType string +} + +// PreprocessMeta 记录缩放与编码结果，供工具输出与排障。 +type PreprocessMeta struct { + OriginalPath string + OriginalBytes int64 + OriginalWidth int + OriginalHeight int + OutputWidth int + OutputHeight int + OutputBytes int + OutputMIMEType string + JPEGQuality int // 0 表示未 JPEG 重编码（原图直传） + PreprocessMode string // passthrough | jpeg +} + +// PreprocessOptions 图片预处理参数。 +type PreprocessOptions struct { + MaxImageBytes int64 + MaxDimension int + JPEGQuality int + MaxPayloadBytes int64 + SkipPreprocessBelowBytes int64 // 0 = 始终压缩；>0 时小图+尺寸合规可直传 +} + +// PreprocessImageFile 读取图片；大图或超尺寸走 imaging 缩放+JPEG，否则可原图直传。 +func PreprocessImageFile(path string, opt PreprocessOptions) (ImagePayload, PreprocessMeta, error) { + var meta PreprocessMeta + meta.OriginalPath = path + + st, err := os.Stat(path) + if err != nil { + return ImagePayload{}, meta, err + } + meta.OriginalBytes = st.Size() + if opt.MaxImageBytes > 0 && st.Size() > opt.MaxImageBytes { + return ImagePayload{}, meta, fmt.Errorf("file size %d exceeds max_image_bytes %d", st.Size(), opt.MaxImageBytes) + } + + cfgW, cfgH, format, err := imageDimensions(path) + if err != nil { + return ImagePayload{}, meta, err + } + meta.OriginalWidth = cfgW + meta.OriginalHeight = cfgH + + maxDim := opt.MaxDimension + if maxDim <= 0 { + maxDim = 2048 + } + maxPayload := opt.MaxPayloadBytes + if maxPayload <= 0 { + maxPayload = 512 * 1024 + } + + if payload, meta, ok, err := tryPassthrough(path, st.Size(), cfgW, cfgH, format, opt, maxDim, maxPayload); ok { + return payload, meta, err + } + + return compressWithImaging(path, opt, maxDim, maxPayload, meta) +} + +func tryPassthrough(path string, size int64, w, h int, format string, opt PreprocessOptions, maxDim int, maxPayload int64) (ImagePayload, PreprocessMeta, bool, error) { + var meta PreprocessMeta + meta.OriginalPath = path + meta.OriginalBytes = size + meta.OriginalWidth = w + meta.OriginalHeight = h + + threshold := opt.SkipPreprocessBelowBytes + if threshold <= 0 { + return ImagePayload{}, meta, false, nil + } + if size > threshold { + return ImagePayload{}, meta, false, nil + } + longEdge := w + if h > longEdge { + longEdge = h + } + if longEdge > maxDim { + return ImagePayload{}, meta, false, nil + } + if size > maxPayload { + return ImagePayload{}, meta, false, nil + } + + raw, err := os.ReadFile(path) + if err != nil { + return ImagePayload{}, meta, false, err + } + mime := mimeFromImageFormat(format) + if mime == "" { + return ImagePayload{}, meta, false, nil + } + + meta.OutputWidth = w + meta.OutputHeight = h + meta.OutputBytes = len(raw) + meta.OutputMIMEType = mime + meta.PreprocessMode = "passthrough" + return ImagePayload{Bytes: raw, MIMEType: mime}, meta, true, nil +} + +func compressWithImaging(path string, opt PreprocessOptions, maxDim int, maxPayload int64, meta PreprocessMeta) (ImagePayload, PreprocessMeta, error) { + src, err := imaging.Open(path) + if err != nil { + return ImagePayload{}, meta, fmt.Errorf("open image: %w", err) + } + bounds := src.Bounds() + meta.OriginalWidth = bounds.Dx() + meta.OriginalHeight = bounds.Dy() + + dst := imaging.Fit(src, maxDim, maxDim, imaging.Lanczos) + outBounds := dst.Bounds() + meta.OutputWidth = outBounds.Dx() + meta.OutputHeight = outBounds.Dy() + + quality := opt.JPEGQuality + if quality <= 0 || quality > 100 { + quality = 82 + } + + dim := maxDim + for attempt := 0; attempt < 6; attempt++ { + if attempt > 0 { + dim = int(float64(dim) * 0.85) + if dim < 256 { + dim = 256 + } + dst = imaging.Fit(src, dim, dim, imaging.Lanczos) + outBounds = dst.Bounds() + meta.OutputWidth = outBounds.Dx() + meta.OutputHeight = outBounds.Dy() + } + q := quality + for q >= 60 { + var buf bytes.Buffer + if err := imaging.Encode(&buf, dst, imaging.JPEG, imaging.JPEGQuality(q)); err != nil { + return ImagePayload{}, meta, fmt.Errorf("encode jpeg: %w", err) + } + if int64(buf.Len()) <= maxPayload { + meta.JPEGQuality = q + meta.OutputBytes = buf.Len() + meta.OutputMIMEType = "image/jpeg" + meta.PreprocessMode = "jpeg" + return ImagePayload{Bytes: buf.Bytes(), MIMEType: "image/jpeg"}, meta, nil + } + q -= 5 + } + quality = 75 + } + return ImagePayload{}, meta, fmt.Errorf("could not compress image under max_payload_bytes %d", maxPayload) +} + +func imageDimensions(path string) (w, h int, format string, err error) { + f, err := os.Open(path) + if err != nil { + return 0, 0, "", err + } + defer f.Close() + cfg, format, err := image.DecodeConfig(f) + if err != nil { + return 0, 0, "", fmt.Errorf("decode image config: %w", err) + } + return cfg.Width, cfg.Height, format, nil +} + +func mimeFromImageFormat(format string) string { + switch strings.ToLower(strings.TrimSpace(format)) { + case "jpeg", "jpg": + return "image/jpeg" + case "png": + return "image/png" + case "gif": + return "image/gif" + case "webp": + return "image/webp" + case "bmp": + return "image/bmp" + case "tiff": + return "image/tiff" + default: + return "" + } +} + +// DecodeImageConfig 用于测试：确认文件可被解码。 +func DecodeImageConfig(path string) (image.Config, string, error) { + f, err := os.Open(path) + if err != nil { + return image.Config{}, "", err + } + defer f.Close() + return image.DecodeConfig(f) +} diff --git a/internal/vision/preprocess_test.go b/internal/vision/preprocess_test.go new file mode 100644 index 00000000..a9b9e068 --- /dev/null +++ b/internal/vision/preprocess_test.go @@ -0,0 +1,109 @@ +package vision + +import ( + "image" + "image/color" + "image/png" + "os" + "path/filepath" + "testing" + + "github.com/disintegration/imaging" +) + +func TestPreprocessImageFile_scalesAndLimitsPayload(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "big.png") + img := imaging.New(3000, 2000, color.White) + if err := imaging.Save(img, path); err != nil { + t.Fatal(err) + } + + out, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxImageBytes: 10 * 1024 * 1024, + MaxDimension: 1024, + JPEGQuality: 85, + MaxPayloadBytes: 600 * 1024, + SkipPreprocessBelowBytes: 0, + }) + if err != nil { + t.Fatal(err) + } + if len(out.Bytes) == 0 { + t.Fatal("empty output") + } + if meta.PreprocessMode != "jpeg" { + t.Fatalf("mode: %s", meta.PreprocessMode) + } + if meta.OutputWidth > 1024 || meta.OutputHeight > 1024 { + t.Fatalf("expected fit within 1024, got %dx%d", meta.OutputWidth, meta.OutputHeight) + } + if int64(len(out.Bytes)) > 600*1024 { + t.Fatalf("payload %d exceeds max", len(out.Bytes)) + } +} + +func TestPreprocessImageFile_passthroughSmallPNG(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "small.png") + if err := imaging.Save(imaging.New(400, 300, color.White), path); err != nil { + t.Fatal(err) + } + + out, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxImageBytes: 5 * 1024 * 1024, + MaxDimension: 2048, + MaxPayloadBytes: 512 * 1024, + SkipPreprocessBelowBytes: 2 * 1024 * 1024, + }) + if err != nil { + t.Fatal(err) + } + if meta.PreprocessMode != "passthrough" { + t.Fatalf("expected passthrough, got %s", meta.PreprocessMode) + } + if out.MIMEType != "image/png" { + t.Fatalf("mime: %s", out.MIMEType) + } + if meta.OutputWidth != 400 || meta.OutputHeight != 300 { + t.Fatalf("dims: %dx%d", meta.OutputWidth, meta.OutputHeight) + } +} + +func TestPreprocessImageFile_passthroughDisabled(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "small.png") + if err := imaging.Save(imaging.New(100, 100, color.White), path); err != nil { + t.Fatal(err) + } + + _, meta, err := PreprocessImageFile(path, PreprocessOptions{ + MaxDimension: 2048, + MaxPayloadBytes: 512 * 1024, + SkipPreprocessBelowBytes: 0, + }) + if err != nil { + t.Fatal(err) + } + if meta.PreprocessMode != "jpeg" { + t.Fatalf("expected jpeg compress, got %s", meta.PreprocessMode) + } +} + +func TestPreprocessImageFile_rejectsOversizeFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "tiny.png") + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + if err := png.Encode(f, image.NewRGBA(image.Rect(0, 0, 2, 2))); err != nil { + t.Fatal(err) + } + f.Close() + + _, _, err = PreprocessImageFile(path, PreprocessOptions{MaxImageBytes: 1}) + if err == nil { + t.Fatal("expected error when file exceeds max_image_bytes") + } +} diff --git a/internal/vision/tool.go b/internal/vision/tool.go new file mode 100644 index 00000000..d31ca928 --- /dev/null +++ b/internal/vision/tool.go @@ -0,0 +1,125 @@ +package vision + +import ( + "context" + "fmt" + "os" + "strings" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +// RegisterAnalyzeImageTool 在 vision.enabled 且 model 已配置时注册 MCP 工具 analyze_image。 +func RegisterAnalyzeImageTool(mcpServer *mcp.Server, cfg *config.Config, logger *zap.Logger) { + if mcpServer == nil || cfg == nil { + return + } + if !cfg.Vision.Ready() { + if cfg.Vision.Enabled && logger != nil { + logger.Warn("vision.enabled 但 vision.model 为空，跳过注册 analyze_image") + } + return + } + + cwd, err := os.Getwd() + if err != nil { + if logger != nil { + logger.Warn("vision: getwd failed, skip analyze_image", zap.Error(err)) + } + return + } + + preOpt := PreprocessOptions{ + MaxImageBytes: cfg.Vision.MaxImageBytesEffective(), + MaxDimension: cfg.Vision.MaxDimensionEffective(), + JPEGQuality: cfg.Vision.JPEGQualityEffective(), + MaxPayloadBytes: cfg.Vision.MaxPayloadBytesEffective(), + SkipPreprocessBelowBytes: cfg.Vision.SkipPreprocessBelowBytesEffective(), + } + client := NewClient(cfg.Vision, cfg.OpenAI) + + tool := mcp.Tool{ + Name: builtin.ToolAnalyzeImage, + Description: "分析服务器上的本地图片并返回文字描述（验证码、UI 元素、报错、架构图要点等）。" + + "输入为文件路径（如用户上传的 chat_uploads 路径或工具截图路径）。" + + "输出仅为文本，不含图片数据。不要对二进制图片使用 read_file 指望理解内容。", + ShortDescription: "分析本地图片并返回文字描述（验证码/UI/报错等）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "图片绝对路径或相对于进程工作目录的路径", + }, + "question": map[string]interface{}{ + "type": "string", + "description": "可选：希望模型重点回答的问题。验证码图建议：只输出验证码字符，不要空格和解释", + }, + }, + "required": []string{"path"}, + }, + } + + handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + path, _ := args["path"].(string) + question, _ := args["question"].(string) + + abs, err := ResolveImagePath(path, cwd) + if err != nil { + return textResult(fmt.Sprintf("路径校验失败: %v", err), true), nil + } + + img, meta, err := PreprocessImageFile(abs, preOpt) + if err != nil { + return textResult(fmt.Sprintf("图片预处理失败: %v", err), true), nil + } + + summary, err := client.Analyze(ctx, img, question) + if err != nil { + return textResult(fmt.Sprintf("视觉模型调用失败: %v", err), true), nil + } + + body := formatAnalysisResult(abs, meta, summary) + return textResult(body, false), nil + } + + mcpServer.RegisterTool(tool, handler) + if logger != nil { + logger.Debug("vision: analyze_image 工具已注册", zap.String("model", cfg.Vision.Model)) + } +} + +func textResult(text string, isError bool) *mcp.ToolResult { + return &mcp.ToolResult{ + Content: []mcp.Content{{Type: "text", Text: text}}, + IsError: isError, + } +} + +func formatAnalysisResult(path string, meta PreprocessMeta, summary string) string { + var b strings.Builder + b.WriteString("## Image analysis\n") + b.WriteString("- **path**: ") + b.WriteString(path) + b.WriteString("\n") + switch meta.PreprocessMode { + case "passthrough": + b.WriteString(fmt.Sprintf("- **preprocess**: passthrough %dx%d, %s, %dKB (original %dKB)\n\n", + meta.OutputWidth, meta.OutputHeight, meta.OutputMIMEType, + (meta.OutputBytes+1023)/1024, (meta.OriginalBytes+1023)/1024)) + default: + b.WriteString(fmt.Sprintf("- **preprocess**: %dx%d → %dx%d, jpeg q=%d, %dKB (original %dKB)\n\n", + meta.OriginalWidth, meta.OriginalHeight, + meta.OutputWidth, meta.OutputHeight, + meta.JPEGQuality, (meta.OutputBytes+1023)/1024, + (meta.OriginalBytes+1023)/1024)) + } + b.WriteString("### Summary\n") + b.WriteString(strings.TrimSpace(summary)) + b.WriteString("\n") + return b.String() +}