package helper import ( "encoding/json" "fmt" "os" "path/filepath" "regexp" "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 仅为前部片段。", } mb, err := json.MarshalIndent(meta, "", " ") if err != nil { return mcp.NewToolResultText(string(b)) } return mcp.NewToolResultText(string(mb)) } // 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 }