FranklinWH-Toolbox/service/cmd/mcp/tools/data/redis.go

package data

import (
	"context"
	"fmt"
	"strings"

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

	"fwh-toolbox-server/cache/redisx"
	"fwh-toolbox-server/cmd/mcp/tools/helper"
)

func registerRedisTools(s *server.MCPServer) {
	s.AddTool(
		mcp.NewTool("data_redis_get",
			mcp.WithDescription("读取 Redis 中指定 key 的值（只读，仅支持 string 和 hash 类型）"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis key")),
			mcp.WithString("type", mcp.Description("key 的数据类型：string（默认）、hash、list、set")),
			mcp.WithNumber("limit", mcp.Description("list/set 类型返回的最大元素数，默认100，最大1000。string/hash 忽略此参数。")),
		),
		handleRedisGet,
	)

	s.AddTool(
		mcp.NewTool("data_redis_exists",
			mcp.WithDescription("检查 Redis key 是否存在，并返回 TTL 信息"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis key")),
		),
		handleRedisExists,
	)
}

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

	keyType := req.GetString("type", "string")

	// limit 仅对 list/set 生效，默认100，上限1000
	limit := int(req.GetFloat("limit", 100))
	if limit < 1 {
		limit = 100
	}
	if limit > 1000 {
		limit = 1000
	}

	switch keyType {
	case "string":
		val, err := redisx.Get(ctx, key)
		if err != nil {
			return helper.EmptyResult(fmt.Sprintf("key 不存在或读取失败: %s (%v)", key, err)), nil
		}
		return helper.JSONResultOrFileAuto(map[string]any{
			"key":   key,
			"type":  "string",
			"value": val,
		}), nil

	case "hash":
		val, err := redisx.HGetAll(ctx, key)
		if err != nil {
			return helper.ServiceError("读取Redis hash", err), nil
		}
		if len(val) == 0 {
			return helper.EmptyResult(fmt.Sprintf("key 不存在或为空: %s", key)), nil
		}
		return helper.JSONResultOrFileAuto(map[string]any{
			"key":   key,
			"type":  "hash",
			"value": val,
		}), nil

	case "list":
		val, err := redisx.LRange(ctx, key, 0, int64(limit-1))
		if err != nil {
			return helper.ServiceError("读取Redis list", err), nil
		}
		if len(val) == 0 {
			return helper.EmptyResult(fmt.Sprintf("key 不存在或为空: %s", key)), nil
		}
		return helper.JSONResultOrFileAuto(map[string]any{
			"key":       key,
			"type":      "list",
			"count":     len(val),
			"limit":     limit,
			"truncated": len(val) == limit,
			"value":     val,
		}), nil

	case "set":
		val, err := redisx.SMembers(ctx, key)
		if err != nil {
			return helper.ServiceError("读取Redis set", err), nil
		}
		if len(val) == 0 {
			return helper.EmptyResult(fmt.Sprintf("key 不存在或为空: %s", key)), nil
		}
		total := len(val)
		truncated := false
		if total > limit {
			val = val[:limit]
			truncated = true
		}
		return helper.JSONResultOrFileAuto(map[string]any{
			"key":       key,
			"type":      "set",
			"total":     total,
			"count":     len(val),
			"limit":     limit,
			"truncated": truncated,
			"value":     val,
		}), nil

	default:
		return helper.ParamInvalidError("type", "可选值: string, hash, list, set"), nil
	}
}

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

	// 安全检查：不允许通配符
	if strings.Contains(key, "*") {
		return helper.ParamInvalidError("key", "不支持通配符，请提供精确 key"), nil
	}

	exists, err := redisx.Exists(ctx, key)
	if err != nil {
		return helper.ServiceError("检查Redis key", err), nil
	}

	if !exists {
		return helper.EmptyResult(fmt.Sprintf("key 不存在: %s", key)), nil
	}

	ttl, err := redisx.TTL(ctx, key)
	if err != nil {
		return helper.ServiceError("获取TTL", err), nil
	}

	ttlStr := "永不过期"
	if ttl > 0 {
		ttlStr = ttl.String()
	} else if ttl == -2 {
		ttlStr = "key 不存在"
	}

	return helper.JSONResultOrFileAuto(map[string]any{
		"key":    key,
		"exists": true,
		"ttl":    ttlStr,
	}), nil
}