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"
)

// 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
}