package helper import ( "bufio" "encoding/json" "fmt" "os" "path/filepath" "regexp" "sort" "strings" "time" "github.com/mark3labs/mcp-go/mcp" ) // AgentBaseDir 是 MCP Agent 允许读写的本地根目录。 // 输入侧 source_path 校验、输出侧结果落盘都限制在此目录及其子目录内。 const AgentBaseDir = "/tmp/fwh-toolbox-agent" // spillDir 大结果落盘目录(AgentBaseDir 的子目录) var spillDir = filepath.Join(AgentBaseDir, "results") // InlineMaxBytes 结果内联返回的最大字节数,超过则落盘到文件。 // 约 32KB,兼顾避免烧 token 与减少不必要的"读文件"往返。 const InlineMaxBytes = 32 * 1024 // previewMaxBytes 落盘时返回给 Agent 的预览内容最大字节数 const previewMaxBytes = 2000 // hintSanitizer 用于清洗文件名 hint,只保留安全字符 var hintSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_\-]+`) // JSONResultOrFile 将数据序列化为 JSON,并在序列化前完成 PII 脱敏(A+B): // - 命中已知强类型:走源头结构体脱敏(优化 A),零额外内存开销。 // - 其余类型:序列化后走流式 token 脱敏(优化 B)。 // - 脱敏失败:返回错误,拒绝输出(fail-closed,绝不泄露原始数据)。 // // 脱敏完成后: // - 小结果(<= InlineMaxBytes)直接内联返回文本 // - 大结果写入本地文件,返回包含 localPath + 预览 + 元信息的摘要 // // hint 用于生成可读的文件名前缀(如 "plane_work_items")。 func JSONResultOrFile(data any, hint string) *mcp.CallToolResult { b, err := marshalSanitized(data) if err != nil { return mcp.NewToolResultError(err.Error()) } return BytesResultOrFile(b, hint) } // marshalSanitized 执行 A+B 脱敏并返回脱敏后的 JSON 字节。 // - 优化 A:命中强类型时先对副本做字段级脱敏,再直接序列化(无需二次处理)。 // - 优化 B:未命中强类型时序列化后做流式 token 脱敏。 func marshalSanitized(data any) ([]byte, error) { if sanitized, ok := sanitizeStrong(data); ok { b, err := json.MarshalIndent(sanitized, "", " ") if err != nil { return nil, fmt.Errorf("序列化失败: %v", err) } return b, nil } b, err := json.MarshalIndent(data, "", " ") if err != nil { return nil, fmt.Errorf("序列化失败: %v", err) } masked, err := MaskJSONBytes(b) if err != nil { return nil, err } return indentJSON(masked), nil } // JSONResultOrFileAuto 与 JSONResultOrFile 相同,但使用通用文件名前缀, // 便于对现有 helper.JSONResult 调用点做无侵入替换。 func JSONResultOrFileAuto(data any) *mcp.CallToolResult { return JSONResultOrFile(data, "result") } // TextResultOrFile 针对可能较大的纯文本返回值做落盘判断: // - 小文本直接内联返回 // - 大文本写入本地文件,返回 localPath + 预览 func TextResultOrFile(text string, hint string) *mcp.CallToolResult { return BytesResultOrFile([]byte(text), hint) } // BytesResultOrFileMasked 针对上游返回的原始 JSON 字节(如 Plane API 响应) // 先做流式 PII 脱敏(优化 B),再走落盘判断。 // - 脱敏失败(含无法解析的非法 JSON):返回错误,拒绝输出(fail-closed)。 // // 相比"unmarshal 成 any 再 JSONResultOrFile",此路径省去一次反序列化+序列化往返。 func BytesResultOrFileMasked(b []byte, hint string) *mcp.CallToolResult { masked, err := MaskJSONBytes(b) if err != nil { return mcp.NewToolResultError(err.Error()) } return BytesResultOrFile(indentJSON(masked), hint) } // BytesResultOrFile 针对已经序列化好的 JSON/文本 bytes 做落盘判断。 func BytesResultOrFile(b []byte, hint string) *mcp.CallToolResult { if len(b) <= InlineMaxBytes { return mcp.NewToolResultText(string(b)) } path, werr := writeSpillFile(hint, b) if werr != nil { // 落盘失败则退化为直接内联返回,保证功能可用(仅牺牲上下文占用) return mcp.NewToolResultText(string(b)) } meta := map[string]any{ "spilled": true, "localPath": path, "bytes": len(b), "preview": truncate(string(b), previewMaxBytes), "note": "结果较大已写入本地文件,请使用文件读取工具按 localPath 读取完整内容;preview 仅为前部片段。", } // 附带结构化摘要:帮助 Agent 在不读回完整文件的情况下了解数据规模与字段,节省 token。 if summary := summarizeJSON(b); summary != nil { meta["summary"] = summary } mb, err := json.MarshalIndent(meta, "", " ") if err != nil { return mcp.NewToolResultText(string(b)) } return mcp.NewToolResultText(string(mb)) } // ExportToFile 用于"全量导出"场景:无论数据大小都落盘为文件,绝不内联。 // 先做 A+B PII 脱敏(与其它返回路径一致,fail-closed),再写入落盘目录, // 返回始终包含 localPath + summary + bytes 的结构化结果(不含数据正文), // 供 Agent 按 localPath 读取完整导出内容。 // // 落盘失败时返回错误(导出场景不做内联退化,避免大结果灌爆上下文)。 func ExportToFile(data any, hint string) *mcp.CallToolResult { b, err := marshalSanitized(data) if err != nil { return mcp.NewToolResultError(err.Error()) } path, werr := writeSpillFile(hint, b) if werr != nil { return mcp.NewToolResultError(fmt.Sprintf("导出落盘失败: %v", werr)) } meta := map[string]any{ "exported": true, "localPath": path, "bytes": len(b), "note": "全量导出结果已写入本地文件,请使用文件读取工具按 localPath 读取完整内容。", } if summary := summarizeJSON(b); summary != nil { meta["summary"] = summary } mb, err := json.MarshalIndent(meta, "", " ") if err != nil { return mcp.NewToolResultError(fmt.Sprintf("导出元信息序列化失败: %v", err)) } return mcp.NewToolResultText(string(mb)) } // marshalSanitizedCompact 对单条数据做 A+B 脱敏并返回紧凑(无缩进/无换行)的 JSON 字节, // 适用于 NDJSON 逐行写入。 // - 强类型命中:sanitizeStrong 后 json.Marshal(紧凑)。 // - 否则:json.Marshal 后 MaskJSONBytes(其输出本身即紧凑)。 func marshalSanitizedCompact(data any) ([]byte, error) { if sanitized, ok := sanitizeStrong(data); ok { b, err := json.Marshal(sanitized) if err != nil { return nil, fmt.Errorf("序列化失败: %v", err) } return b, nil } b, err := json.Marshal(data) if err != nil { return nil, fmt.Errorf("序列化失败: %v", err) } return MaskJSONBytes(b) } // ExportStreamNDJSON 流式导出:逐条消费 next(),边脱敏边写入 NDJSON 文件(每行一个 JSON 对象), // 内存占用为常数级(任意时刻仅持有一条记录),适合超大数据量导出。 // // next 语义:返回 (记录, more, err)。 // - more=false 且 err=nil:正常结束。 // - err!=nil:迭代出错,导出中止并清理半成品文件。 // // 任一条脱敏/写入失败:删除半成品文件并返回错误(fail-closed,绝不留下未脱敏残片)。 // 成功返回 {exported:true, format:"ndjson", localPath, count, summary}(不含数据正文)。 func ExportStreamNDJSON(next func() (any, bool, error), hint string) *mcp.CallToolResult { if err := os.MkdirAll(spillDir, 0755); err != nil { return mcp.NewToolResultError(fmt.Sprintf("创建导出目录失败: %v", err)) } name := fmt.Sprintf("%s_%d.ndjson", sanitizeHint(hint), time.Now().UnixNano()) path := filepath.Join(spillDir, name) f, err := os.Create(path) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("创建导出文件失败: %v", err)) } // 失败时清理半成品文件的辅助闭包 failClose := func(msg string) *mcp.CallToolResult { f.Close() _ = os.Remove(path) return mcp.NewToolResultError(msg) } w := bufio.NewWriter(f) var count int64 var fields []string // 采样首条记录的字段名用于 summary for { rec, more, ierr := next() if ierr != nil { return failClose(fmt.Sprintf("导出迭代失败: %v", ierr)) } if !more { break } line, merr := marshalSanitizedCompact(rec) if merr != nil { return failClose(merr.Error()) } // 采样首条字段名(尽力而为,失败不影响导出) if count == 0 { var m map[string]any if json.Unmarshal(line, &m) == nil { fields = sortedKeys(m) } } if _, werr := w.Write(line); werr != nil { return failClose(fmt.Sprintf("写入导出文件失败: %v", werr)) } if werr := w.WriteByte('\n'); werr != nil { return failClose(fmt.Sprintf("写入导出文件失败: %v", werr)) } count++ } if werr := w.Flush(); werr != nil { return failClose(fmt.Sprintf("刷新导出文件失败: %v", werr)) } if cerr := f.Close(); cerr != nil { _ = os.Remove(path) return mcp.NewToolResultError(fmt.Sprintf("关闭导出文件失败: %v", cerr)) } if count == 0 { _ = os.Remove(path) return EmptyResult("符合条件的记录为 0,未生成导出文件") } meta := map[string]any{ "exported": true, "format": "ndjson", "localPath": path, "count": count, "note": "全量导出结果已写入 NDJSON 文件(每行一个 JSON 对象),请使用文件读取工具按 localPath 逐行读取解析。", } if len(fields) > 0 { meta["summary"] = map[string]any{ "topType": "ndjson", "count": count, "itemType": "object", "fields": fields, } } mb, err := json.MarshalIndent(meta, "", " ") if err != nil { return mcp.NewToolResultError(fmt.Sprintf("导出元信息序列化失败: %v", err)) } return mcp.NewToolResultText(string(mb)) } // summarizeJSON 尝试从已序列化的 JSON 字节中提取结构化摘要,供 Agent 快速了解 // 落盘数据的规模与结构,从而尽量避免读回完整文件。解析失败时返回 nil(不影响主流程)。 // // 摘要内容: // - topType: 顶层类型(array / object / scalar) // - 顶层为数组:count(元素个数)、itemType、fields(首个对象元素的字段名,采样) // - 顶层为对象:fields(顶层字段名) func summarizeJSON(b []byte) map[string]any { var v any dec := json.NewDecoder(strings.NewReader(string(b))) dec.UseNumber() if err := dec.Decode(&v); err != nil { return nil } switch t := v.(type) { case []any: s := map[string]any{ "topType": "array", "count": len(t), } if len(t) > 0 { if obj, ok := t[0].(map[string]any); ok { s["itemType"] = "object" s["fields"] = sortedKeys(obj) } else { s["itemType"] = scalarKind(t[0]) } } return s case map[string]any: return map[string]any{ "topType": "object", "fields": sortedKeys(t), } default: return map[string]any{"topType": "scalar"} } } // sortedKeys 返回 map 的键列表(稳定排序,便于 Agent 阅读与 diff)。 func sortedKeys(m map[string]any) []string { keys := make([]string, 0, len(m)) for k := range m { keys = append(keys, k) } sort.Strings(keys) return keys } // scalarKind 返回标量的粗粒度类型名。 func scalarKind(v any) string { switch v.(type) { case string: return "string" case bool: return "bool" case nil: return "null" default: // json.Number 及其它 return "number" } } // writeSpillFile 将 bytes 写入落盘目录,返回绝对路径。 func writeSpillFile(hint string, b []byte) (string, error) { if err := os.MkdirAll(spillDir, 0755); err != nil { return "", err } name := fmt.Sprintf("%s_%d.json", sanitizeHint(hint), time.Now().UnixNano()) path := filepath.Join(spillDir, name) if err := os.WriteFile(path, b, 0644); err != nil { return "", err } return path, nil } // sanitizeHint 清洗文件名前缀,避免路径穿越和非法字符 func sanitizeHint(hint string) string { hint = hintSanitizer.ReplaceAllString(hint, "_") hint = strings.Trim(hint, "_") if hint == "" { hint = "result" } if len(hint) > 64 { hint = hint[:64] } return hint } // truncate 截断字符串到指定字节数(按 rune 边界安全截断) func truncate(s string, max int) string { if len(s) <= max { return s } // 从 max 处向前回退到 rune 边界 end := max for end > 0 && !utf8RuneStart(s[end]) { end-- } return s[:end] + "\n...(truncated)" } // utf8RuneStart 判断字节是否为 UTF-8 编码的起始字节 func utf8RuneStart(b byte) bool { // 0xxxxxxx (ASCII) 或 11xxxxxx (多字节起始);10xxxxxx 为续字节 return b&0xC0 != 0x80 } // ValidateSourcePath 校验输入侧的本地文件路径: // - 必须位于 AgentBaseDir(/tmp/fwh-toolbox-agent)目录或其子目录内 // - 清理 ".." 等路径穿越 // // 返回清理后的绝对路径;非法路径返回 error。 func ValidateSourcePath(path string) (string, error) { if path == "" { return "", fmt.Errorf("路径为空") } abs, err := filepath.Abs(filepath.Clean(path)) if err != nil { return "", fmt.Errorf("路径解析失败: %v", err) } base := filepath.Clean(AgentBaseDir) // 允许 base 本身或 base 的子路径;用分隔符前缀避免 /tmp/fwh-toolbox-agent-evil 误通过 if abs != base && !strings.HasPrefix(abs, base+string(os.PathSeparator)) { return "", fmt.Errorf("路径不在允许目录内,仅允许 %s 及其子目录", AgentBaseDir) } return abs, nil }