package helper import ( "bufio" "encoding/json" "fmt" "os" "path/filepath" "regexp" "sort" "strings" "time" "github.com/mark3labs/mcp-go/mcp" "fwh-toolbox-server/utils/agentfiles" ) // AgentBaseDir 是 MCP Agent 允许读写的本地根目录。 // 输入侧 source_path 校验、输出侧结果落盘都限制在此目录及其子目录内。 const AgentBaseDir = "/tmp/fwh-toolbox-agent" // init 断言:本包 AgentBaseDir 必须与 agentfiles.AgentBaseDir 一致, // 否则落盘目录与清理目录不匹配会导致文件永不被清理。不一致时 MCP 启动即 panic, // 让问题在部署时立刻暴露,而非静默漂移。 func init() { if AgentBaseDir != agentfiles.AgentBaseDir { panic(fmt.Sprintf( "AgentBaseDir 不一致: helper=%q agentfiles=%q(两处必须相同)", AgentBaseDir, agentfiles.AgentBaseDir)) } } // spillDir 大结果落盘目录(AgentBaseDir 的子目录) var spillDir = filepath.Join(AgentBaseDir, "results") // hintSanitizer 用于清洗文件名 hint,只保留安全字符 var hintSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_\-]+`) // JSONResultOrFile 将数据序列化为 JSON,并在序列化前完成 PII 脱敏(A+B): // - 命中已知强类型:走源头结构体脱敏(优化 A),零额外内存开销。 // - 其余类型:序列化后走流式 token 脱敏(优化 B)。 // - 脱敏失败:返回错误,拒绝输出(fail-closed,绝不泄露原始数据)。 // // 脱敏完成后无条件落盘(数据类返回统一策略):无论数据大小都写入本地文件, // 返回体只含 localPath + bytes + summary(不含数据正文),避免数据进入模型上下文。 // // 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 + summary(不含正文)。 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 做无条件落盘。 // // 语义(数据类返回):无论数据大小都写入本地文件,返回结果只包含 // localPath + bytes + summary(不含数据正文),供 Agent 按 localPath 读取完整内容。 // 这样任意"带数据"的查询/导出结果都不会进入模型上下文,从根本上避免上下文膨胀。 // // 落盘失败时返回错误(不再退化为内联,避免大结果灌爆上下文)。 func BytesResultOrFile(b []byte, hint string) *mcp.CallToolResult { return spillResult(b, hint) } // spillResult 是"带数据结果一律落盘"的统一实现:写盘 + 组装 localPath/summary 元信息。 // 供 BytesResultOrFile / ExportToFile 等数据类返回路径复用。 func spillResult(b []byte, hint string) *mcp.CallToolResult { path, werr := writeSpillFile(hint, b) if werr != nil { return mcp.NewToolResultError(fmt.Sprintf("结果落盘失败: %v", werr)) } meta := map[string]any{ "localPath": path, "bytes": len(b), "note": "结果数据已写入本地文件,请使用文件读取工具按 localPath 读取完整内容。", } // 附带结构化摘要:帮助 Agent 在不读回完整文件的情况下了解数据规模与字段,节省 token。 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)) } // ExportToFile 用于"全量导出"场景:将数据一次性序列化后落盘。 // 先做 A+B PII 脱敏(与其它返回路径一致,fail-closed),再走统一的落盘返回, // 结果只包含 localPath + bytes + summary(不含数据正文)。 // // 注:数据类返回现已统一为无条件落盘,本函数与 JSONResultOrFile 行为一致, // 保留独立入口以表达"全量导出"语义、便于调用点阅读。 func ExportToFile(data any, hint string) *mcp.CallToolResult { b, err := marshalSanitized(data) if err != nil { return mcp.NewToolResultError(err.Error()) } return spillResult(b, hint) } // 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 } // 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 }