FranklinWH-Toolbox/service/cmd/mcp/tools/storage/file.go

package storage

import (
	"context"
	"encoding/base64"
	"fmt"
	"path/filepath"

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

	"fwh-toolbox-server/cmd/mcp/tools/helper"
	"fwh-toolbox-server/utils/filestore"
)

// mcpFileStore MCP 专用的文件存储实例
var mcpFileStore *filestore.FileStore

// InitFileStore 初始化 MCP FileStore（由 main.go 调用）
func InitFileStore(fs *filestore.FileStore) {
	mcpFileStore = fs
}

// GetFileStore 获取 MCP FileStore 实例
func GetFileStore() *filestore.FileStore {
	return mcpFileStore
}

func registerFileTools(s *server.MCPServer) {
	s.AddTool(
		mcp.NewTool("storage_file_read",
			mcp.WithDescription("读取用户上传的文本文件内容（返回纯文本）。用户通过前端上传的文件存储在 /uploads/ 路径下。仅适用于文本类文件。"),
			mcp.WithString("path", mcp.Required(), mcp.Description("文件路径，格式为 /uploads/{filename}")),
		),
		handleFileRead,
	)

	s.AddTool(
		mcp.NewTool("storage_file_read_base64",
			mcp.WithDescription("读取用户上传的文件内容并返回 base64 编码。支持任意文件类型（图片、PDF、压缩包等）。图片文件读取后可直接用于 AI 视觉分析。"),
			mcp.WithString("path", mcp.Required(), mcp.Description("文件路径，格式为 /uploads/{filename}")),
		),
		handleFileReadBase64,
	)

	s.AddTool(
		mcp.NewTool("storage_file_write",
			mcp.WithDescription("将文本内容写入文件存储（MinIO）。可用于保存分析结果、生成报告等。最大 20MB。文件路径会自动添加 /uploads/ 前缀。如需上传二进制文件请使用 storage_file_upload。"),
			mcp.WithString("path", mcp.Required(), mcp.Description("文件路径，格式为 /uploads/{filename}（如 /uploads/report.csv）")),
			mcp.WithString("content", mcp.Required(), mcp.Description("要写入的文本内容（最大 20MB）")),
			mcp.WithString("userId", mcp.Description("用户标识，用于文件路径隔离（尽量填写，不传则使用默认值 mcp-agent）")),
			mcp.WithString("sessionId", mcp.Description("会话标识，用于文件路径隔离（尽量填写，不传则使用默认值 mcp-session）")),
		),
		handleFileWrite,
	)

	s.AddTool(
		mcp.NewTool("storage_file_list",
			mcp.WithDescription("列出当前会话中用户上传的所有文件，返回文件名、路径、大小等信息。"),
		),
		handleFileList,
	)

	s.AddTool(
		mcp.NewTool("storage_file_upload",
			mcp.WithDescription("将任意类型文件上传到文件存储（MinIO）。内容需使用 base64 编码传入。支持二进制文件（图片、PDF、压缩包、Excel 等）和文本文件。最大 100MB。"),
			mcp.WithString("fileName", mcp.Required(), mcp.Description("目标文件名（如 data.zip、report.pdf、image.png）")),
			mcp.WithString("content_base64", mcp.Required(), mcp.Description("文件内容的 base64 编码字符串（可通过读取本地文件后 base64 编码获得）")),
			mcp.WithString("userId", mcp.Description("用户标识，用于文件路径隔离（尽量填写，不传则使用默认值 mcp-agent）")),
			mcp.WithString("sessionId", mcp.Description("会话标识，用于文件路径隔离（尽量填写，不传则使用默认值 mcp-session）")),
		),
		handleFileUpload,
	)
}

func handleFileRead(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	path, err := req.RequireString("path")
	if err != nil {
		return helper.ParamError("path"), nil
	}

	if mcpFileStore == nil {
		return helper.ServiceError("文件存储", fmt.Errorf("未初始化")), nil
	}

	if !mcpFileStore.HasFile(path) {
		return helper.EmptyResult(fmt.Sprintf("文件不存在: %s", path)), nil
	}

	content, err := mcpFileStore.ReadFile(path)
	if err != nil {
		return helper.ServiceError("读取文件", err), nil
	}

	return helper.TextResult(content), nil
}

func handleFileReadBase64(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	path, err := req.RequireString("path")
	if err != nil {
		return helper.ParamError("path"), nil
	}

	if mcpFileStore == nil {
		return helper.ServiceError("文件存储", fmt.Errorf("未初始化")), nil
	}

	if !mcpFileStore.HasFile(path) {
		return helper.EmptyResult(fmt.Sprintf("文件不存在: %s", path)), nil
	}

	data, err := mcpFileStore.ReadFileBytes(path)
	if err != nil {
		return helper.ServiceError("读取文件", err), nil
	}

	encoded := base64.StdEncoding.EncodeToString(data)

	return helper.JSONResult(map[string]any{
		"path":           path,
		"size":           len(data),
		"content_base64": encoded,
	}), nil
}

func handleFileList(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	if mcpFileStore == nil {
		return helper.ServiceError("文件存储", fmt.Errorf("未初始化")), nil
	}

	files := mcpFileStore.GetFiles()
	if len(files) == 0 {
		return helper.EmptyResult("当前没有上传的文件"), nil
	}

	return helper.JSONResult(files), nil
}

func handleFileWrite(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	path, err := req.RequireString("path")
	if err != nil {
		return helper.ParamError("path"), nil
	}
	content, err := req.RequireString("content")
	if err != nil {
		return helper.ParamError("content"), nil
	}

	// 限制文本内容大小（最大 20MB）
	const maxSize = 20 * 1024 * 1024
	if len(content) > maxSize {
		return helper.ParamInvalidError("content", fmt.Sprintf("内容大小 %d bytes 超过 20MB 限制", len(content))), nil
	}

	if mcpFileStore == nil {
		return helper.ServiceError("文件存储", fmt.Errorf("未初始化")), nil
	}

	userId := req.GetString("userId", "mcp-agent")
	sessionId := req.GetString("sessionId", "mcp-session")

	record, err := mcpFileStore.WriteFile(userId, sessionId, path, content)
	if err != nil {
		return helper.ServiceError("写入文件", err), nil
	}

	return helper.JSONResult(map[string]any{
		"message":  "文件写入成功",
		"path":     record.Path,
		"fileName": record.FileName,
		"size":     record.Size,
		"key":      record.Key,
	}), nil
}

func handleFileUpload(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	fileName, err := req.RequireString("fileName")
	if err != nil {
		return helper.ParamError("fileName"), nil
	}
	contentBase64, err := req.RequireString("content_base64")
	if err != nil {
		return helper.ParamError("content_base64"), nil
	}

	if mcpFileStore == nil {
		return helper.ServiceError("文件存储", fmt.Errorf("未初始化")), nil
	}

	// base64 解码
	data, err := base64.StdEncoding.DecodeString(contentBase64)
	if err != nil {
		// 尝试 RawStdEncoding（无 padding）
		data, err = base64.RawStdEncoding.DecodeString(contentBase64)
		if err != nil {
			return helper.ParamInvalidError("content_base64", "base64 解码失败"), nil
		}
	}

	// 限制文件大小（最大 100MB）
	const maxSize = 100 * 1024 * 1024
	if len(data) > maxSize {
		return helper.ParamInvalidError("content_base64", fmt.Sprintf("文件大小 %d bytes 超过 100MB 限制", len(data))), nil
	}

	userId := req.GetString("userId", "mcp-agent")
	sessionId := req.GetString("sessionId", "mcp-session")

	virtualPath := fmt.Sprintf("/uploads/%s", filepath.Base(fileName))

	record, err := mcpFileStore.WriteBytesWithSession(userId, sessionId, virtualPath, data)
	if err != nil {
		return helper.ServiceError("上传文件", err), nil
	}

	return helper.JSONResult(map[string]any{
		"message":  "文件上传成功",
		"path":     record.Path,
		"fileName": record.FileName,
		"size":     record.Size,
		"key":      record.Key,
	}), nil
}