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

package helper

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/mark3labs/mcp-go/mcp"
)

func TestJSONResultOrFile_SmallInline(t *testing.T) {
	res := JSONResultOrFile(map[string]any{"a": 1}, "small")
	if res == nil || len(res.Content) == 0 {
		t.Fatal("empty result")
	}
	txt := textOf(t, res)
	if strings.Contains(txt, "spilled") {
		t.Fatalf("small result should be inline, got: %s", txt)
	}
	if !strings.Contains(txt, "\"a\": 1") {
		t.Fatalf("unexpected inline content: %s", txt)
	}
}

func TestJSONResultOrFile_LargeSpill(t *testing.T) {
	// 构造超过 InlineMaxBytes 的数据
	big := make([]map[string]string, 0, 5000)
	for i := 0; i < 5000; i++ {
		big = append(big, map[string]string{"k": "value-value-value-value"})
	}
	res := JSONResultOrFile(big, "large_test")
	txt := textOf(t, res)

	var meta map[string]any
	if err := json.Unmarshal([]byte(txt), &meta); err != nil {
		t.Fatalf("expected JSON meta, got: %s", txt)
	}
	if meta["spilled"] != true {
		t.Fatalf("expected spilled=true, got: %v", meta["spilled"])
	}
	path, _ := meta["localPath"].(string)
	if path == "" {
		t.Fatal("missing localPath")
	}
	// 落盘路径必须在 AgentBaseDir 内
	if !strings.HasPrefix(path, AgentBaseDir) {
		t.Fatalf("spill path outside base dir: %s", path)
	}
	// 文件应真实存在且内容为完整 JSON
	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("spill file not readable: %v", err)
	}
	var parsed []map[string]string
	if err := json.Unmarshal(data, &parsed); err != nil {
		t.Fatalf("spill file not valid JSON: %v", err)
	}
	if len(parsed) != 5000 {
		t.Fatalf("spill file content incomplete: %d", len(parsed))
	}
	_ = os.Remove(path)
}

func TestValidateSourcePath(t *testing.T) {
	okDir := filepath.Join(AgentBaseDir, "downloads")
	_ = os.MkdirAll(okDir, 0755)

	cases := []struct {
		name    string
		in      string
		wantErr bool
	}{
		{"base itself", AgentBaseDir, false},
		{"subdir file", filepath.Join(AgentBaseDir, "downloads", "a.json"), false},
		{"nested subdir", filepath.Join(AgentBaseDir, "x", "y", "z.txt"), false},
		{"traversal escape", filepath.Join(AgentBaseDir, "..", "etc", "passwd"), true},
		{"outside abs", "/etc/passwd", true},
		{"sibling prefix trick", "/tmp/fwh-toolbox-agent-evil/x", true},
		{"empty", "", true},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			_, err := ValidateSourcePath(c.in)
			if c.wantErr && err == nil {
				t.Fatalf("expected error for %q", c.in)
			}
			if !c.wantErr && err != nil {
				t.Fatalf("unexpected error for %q: %v", c.in, err)
			}
		})
	}
}

func textOf(t *testing.T, res *mcp.CallToolResult) string {
	t.Helper()
	if res == nil || len(res.Content) == 0 {
		t.Fatal("empty result content")
	}
	tc, ok := res.Content[0].(mcp.TextContent)
	if !ok {
		t.Fatalf("content is not text: %T", res.Content[0])
	}
	return tc.Text
}