FranklinWH-Toolbox/service/cmd/mcp/tools/helper/spill.go

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
}