FranklinWH-Toolbox/service/agentmcp/kit/result.go

package kit

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/csv"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"reflect"
	"strings"
	"unicode/utf8"

	"github.com/mark3labs/mcp-go/mcp"
	"go.mongodb.org/mongo-driver/mongo"
	"go.uber.org/zap"

	"fwh-toolbox-server/utils/pii"
)

// MaxResultChars 单次结果（非导出）的最大字符数。低于 AgentServer 网关的 10 万字截断阈值，
// 避免触发"前 10 万字仍进上下文"的路径。
const MaxResultChars = 30000

// MaxExportBytes 导出文件上限。
const MaxExportBytes = 20 << 20

// ---- 错误 ----

// Fail 返回给模型的业务错误（参数不对、无数据等）。
func Fail(err error) *mcp.CallToolResult {
	var ae *ArgError
	if errors.As(err, &ae) {
		return mcp.NewToolResultError(ae.Msg)
	}
	return mcp.NewToolResultError(err.Error())
}

// Failf 以格式化文本返回错误。
func Failf(format string, a ...any) *mcp.CallToolResult {
	return mcp.NewToolResultError(fmt.Sprintf(format, a...))
}

// Upstream 上游（FWH 平台、数据库、Plane 等）调用失败。
func Upstream(source string, err error) *mcp.CallToolResult {
	return mcp.NewToolResultError(UpstreamText(source, err) + "。可以稍后重试，或缩小查询范围。")
}

// UpstreamText 写给模型的上游错误说明：只给数据源、错误类别与错误编号。
// 原始错误可能含内部地址、SQL、库表名，只写入日志（按错误编号检索），不返回给模型。
func UpstreamText(source string, err error) string {
	ref := errorRef()
	zap.L().Warn("[agent-mcp] 上游调用失败", zap.String("source", source), zap.String("ref", ref), zap.Error(err))
	return fmt.Sprintf("调用%s失败（%s，错误编号 %s）", source, errorKind(err), ref)
}

func errorKind(err error) string {
	var ne net.Error
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		return "请求超时"
	case errors.Is(err, context.Canceled):
		return "请求已取消"
	case errors.Is(err, mongo.ErrNoDocuments):
		return "没有找到对应记录"
	case errors.As(err, &ne) && ne.Timeout():
		return "请求超时"
	case errors.As(err, &ne):
		return "网络连接失败"
	}
	return "数据源返回错误"
}

func errorRef() string {
	b := make([]byte, 4)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

// Empty 无数据时的统一返回（不是错误）。
func Empty(msg string) *mcp.CallToolResult {
	if msg == "" {
		msg = "没有查询到数据。"
	}
	return mcp.NewToolResultText(msg)
}

// ---- 对象与列表 ----

// Object 单个对象：脱敏、去掉空值后以紧凑 JSON 返回。
func Object(v any) *mcp.CallToolResult {
	tree, err := sanitize(v)
	if err != nil {
		return mcp.NewToolResultError(err.Error())
	}
	b, _ := json.Marshal(tree)
	text := string(b)
	if utf8.RuneCountInString(text) > MaxResultChars {
		text = truncateRunes(text, MaxResultChars) + "\n（结果过长已截断，请缩小查询范围）"
	}
	return mcp.NewToolResultText(text)
}

// ListResult 列表结果的结构。
type ListResult struct {
	Total     int    `json:"total"`
	Returned  int    `json:"returned"`
	Truncated bool   `json:"truncated"`
	Hint      string `json:"hint,omitempty"`
	Items     []any  `json:"items"`
}

// List 列表：脱敏、去空值，最多返回 limit 条；序列化后超过 MaxResultChars 时继续减少条数。
// items 必须是切片（nil 视为空）；total < 0 时以切片长度作为总数。
func List(items any, total, limit int, hint string) *mcp.CallToolResult {
	tree, err := sanitize(items)
	if err != nil {
		return mcp.NewToolResultError(err.Error())
	}
	arr, _ := tree.([]any)
	if total < 0 {
		total = len(arr)
	}
	if len(arr) == 0 {
		return Empty("")
	}
	r := buildList(arr, total, limit, hint)
	b, _ := json.Marshal(r)
	return mcp.NewToolResultText(string(b))
}

func buildList(arr []any, total, limit int, hint string) ListResult {
	n := len(arr)
	if limit > 0 && n > limit {
		n = limit
	}
	for {
		r := ListResult{Total: total, Returned: n, Truncated: n < total, Items: arr[:n]}
		if r.Truncated {
			r.Hint = hint
			if r.Hint == "" {
				r.Hint = "结果已截断，请缩小查询范围或分页查询"
			}
		}
		b, _ := json.Marshal(r)
		if utf8.RuneCount(b) <= MaxResultChars || n <= 1 {
			return r
		}
		n = n * 3 / 4
	}
}

// Table 行多、字段固定的数据：首行列名的 TSV，比 JSON 省字符。单元格按列名脱敏。
func Table(cols []string, rows [][]string, limit int, hint string) *mcp.CallToolResult {
	if len(rows) == 0 {
		return Empty("")
	}
	total := len(rows)
	n := total
	if limit > 0 && n > limit {
		n = limit
	}
	var buf strings.Builder
	header := fmt.Sprintf("# total=%d", total)
	colLine := strings.Join(cols, "\t")
	buf.WriteString(colLine)
	buf.WriteByte('\n')
	chars := utf8.RuneCountInString(colLine) + 1
	written := 0
	for _, row := range rows[:n] {
		line := tsvLine(cols, row)
		lc := utf8.RuneCountInString(line) + 1
		// 预留 200 字符给头部说明
		if chars+lc+200 > MaxResultChars {
			break
		}
		buf.WriteString(line)
		buf.WriteByte('\n')
		chars += lc
		written++
	}
	header += fmt.Sprintf(" returned=%d truncated=%t", written, written < total)
	if written < total {
		if hint == "" {
			hint = "结果已截断，请缩小查询范围"
		}
		header += "\n# " + hint
	}
	return mcp.NewToolResultText(header + "\n" + buf.String())
}

func tsvLine(cols []string, row []string) string {
	cells := make([]string, len(row))
	for i, c := range row {
		if i < len(cols) {
			c = pii.MaskField(cols[i], c)
		}
		cells[i] = strings.NewReplacer("\t", " ", "\n", " ", "\r", " ").Replace(c)
	}
	return strings.Join(cells, "\t")
}

// Export 导出：CSV 以 EmbeddedResource（text/csv）返回，并附文本摘要。
// AgentServer 网关会把资源内容保存到会话工作区（agent-mcp.md 5.2）。
func Export(uri string, cols []string, rows [][]string, summary string) *mcp.CallToolResult {
	if len(rows) == 0 {
		return Empty("没有可导出的数据。")
	}
	var buf bytes.Buffer
	w := csv.NewWriter(&buf)
	_ = w.Write(cols)
	for _, row := range rows {
		masked := make([]string, len(row))
		for i, c := range row {
			if i < len(cols) {
				c = pii.MaskField(cols[i], c)
			}
			masked[i] = c
		}
		_ = w.Write(masked)
		if buf.Len() > MaxExportBytes {
			return Failf("导出数据超过 %d MB，请缩小时间范围后分段导出", MaxExportBytes>>20)
		}
	}
	w.Flush()

	preview := rows
	if len(preview) > 5 {
		preview = preview[:5]
	}
	var head strings.Builder
	fmt.Fprintf(&head, "已导出 %d 行，%d 列，%.1f KB。\n资源：%s\n列：%s\n", len(rows), len(cols), float64(buf.Len())/1024, uri, strings.Join(cols, ", "))
	if summary != "" {
		head.WriteString(summary + "\n")
	}
	head.WriteString("前 5 行：\n")
	head.WriteString(strings.Join(cols, "\t") + "\n")
	for _, r := range preview {
		head.WriteString(tsvLine(cols, r) + "\n")
	}
	head.WriteString("完整数据在资源文件中，请用 Python（pandas）读取分析，不要把全文读入上下文。")

	return &mcp.CallToolResult{
		Content: []mcp.Content{
			mcp.NewTextContent(head.String()),
			mcp.NewEmbeddedResource(mcp.TextResourceContents{URI: uri, MIMEType: "text/csv", Text: buf.String()}),
		},
	}
}

// ---- 内部 ----

// sanitize 序列化 → 脱敏 → 解析为通用树 → 去掉空值。
func sanitize(v any) (any, error) {
	if v == nil || (reflect.ValueOf(v).Kind() == reflect.Ptr && reflect.ValueOf(v).IsNil()) {
		return nil, nil
	}
	b, err := pii.Marshal(v)
	if err != nil {
		return nil, err
	}
	dec := json.NewDecoder(bytes.NewReader(b))
	dec.UseNumber()
	var tree any
	if err := dec.Decode(&tree); err != nil {
		return nil, fmt.Errorf("结果解析失败: %w", err)
	}
	return Prune(tree), nil
}

// Prune 去掉 null、空字符串、空数组、空对象（递归）。数组中的元素保留位置，只清理其内部。
func Prune(v any) any {
	switch t := v.(type) {
	case map[string]any:
		for k, val := range t {
			p := Prune(val)
			if isEmpty(p) {
				delete(t, k)
			} else {
				t[k] = p
			}
		}
		return t
	case []any:
		for i := range t {
			t[i] = Prune(t[i])
		}
		return t
	default:
		return v
	}
}

func isEmpty(v any) bool {
	switch t := v.(type) {
	case nil:
		return true
	case string:
		return t == ""
	case map[string]any:
		return len(t) == 0
	case []any:
		return len(t) == 0
	}
	return false
}

func truncateRunes(s string, n int) string {
	if utf8.RuneCountInString(s) <= n {
		return s
	}
	r := []rune(s)
	return string(r[:n])
}