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：
//   - 小结果（<= InlineMaxBytes）直接内联返回文本
//   - 大结果写入本地文件，返回包含 localPath + 预览 + 元信息的摘要
//
// hint 用于生成可读的文件名前缀（如 "plane_work_items"）。
func JSONResultOrFile(data any, hint string) *mcp.CallToolResult {
	b, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("序列化失败: %v", err))
	}
	return BytesResultOrFile(b, hint)
}

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

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