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

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
}