mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-19 01:17:16 +02:00
Add files via upload
This commit is contained in:
公明
@@ -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
|
||||
}
|
||||
|
||||
// 以下所有方法已不再使用，已删除以简化代码
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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]) + "…"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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...)
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
// <persisted-output> 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
|
||||
}
|
||||
@@ -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, "<persisted-output>") || !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, "<persisted-output>") {
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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() }
|
||||
@@ -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 拼长链。`
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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 <persisted-output> 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(
|
||||
"<persisted-output>\nOutput too large (%d). Full output saved to: %s\nUse %s to read.\n</persisted-output>",
|
||||
originalSize, filePath, readFileHint,
|
||||
)
|
||||
if len(minimal) > maxBytes {
|
||||
core := fmt.Sprintf("<persisted-output>Full output saved to: %s</persisted-output>", filePath)
|
||||
if len(core) <= maxBytes {
|
||||
return core
|
||||
}
|
||||
// Path longer than budget: keep as much of the path as possible after a short prefix.
|
||||
prefix := "<persisted-output>Full output saved to: "
|
||||
suffix := "</persisted-output>"
|
||||
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(
|
||||
"<persisted-output>\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</persisted-output>",
|
||||
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
|
||||
}
|
||||
@@ -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, "<persisted-output>") {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user