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"
)

// Redis 只读工具集。
// 说明：本工具集仅提供只读命令，不包含任何写操作。数据均位于业务约定的 Redis DB 12
//（业务缓存 + asynq 任务队列），由 MCP 进程启动时的 cache.Init 固定连接。

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

	s.AddTool(
		mcp.NewTool("data_redis_mget",
			mcp.WithDescription("批量读取多个 string key 的值（只读）。"),
			mcp.WithString("keys", mcp.Required(), mcp.Description("多个 key，用英文逗号分隔")),
		),
		handleRedisMGet,
	)

	s.AddTool(
		mcp.NewTool("data_redis_hget",
			mcp.WithDescription("读取 hash 中指定 field 的值（只读）。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis hash key")),
			mcp.WithString("field", mcp.Required(), mcp.Description("hash field 名")),
		),
		handleRedisHGet,
	)

	s.AddTool(
		mcp.NewTool("data_redis_type",
			mcp.WithDescription("查看 key 的类型及规模概览（只读）：类型、TTL、以及对应类型的长度/元素数。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis key")),
		),
		handleRedisType,
	)

	s.AddTool(
		mcp.NewTool("data_redis_exists",
			mcp.WithDescription("检查 Redis key 是否存在，并返回 TTL 信息（只读）。key 支持精确匹配。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis key")),
		),
		handleRedisExists,
	)

	s.AddTool(
		mcp.NewTool("data_redis_ttl",
			mcp.WithDescription("查看 key 的存活时间 TTL（只读）。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Redis key")),
		),
		handleRedisTTL,
	)

	s.AddTool(
		mcp.NewTool("data_redis_scan",
			mcp.WithDescription("游标遍历 key（只读，非阻塞，推荐用于扫描）。返回本批 keys 与下一个游标（cursor=0 表示结束）。"),
			mcp.WithString("match", mcp.Description("匹配模式，如 user:* ；为空则匹配全部")),
			mcp.WithNumber("cursor", mcp.Description("起始游标，首次传 0，后续传上次返回的 next_cursor")),
			mcp.WithNumber("count", mcp.Description("每批返回数量提示，默认100，最大1000")),
		),
		handleRedisScan,
	)

	s.AddTool(
		mcp.NewTool("data_redis_keys",
			mcp.WithDescription("按 pattern 返回匹配的 key（只读）。注意：KEYS 为阻塞命令，大库慎用，建议优先用 data_redis_scan。"),
			mcp.WithString("pattern", mcp.Required(), mcp.Description("匹配模式，如 user:* 或 * ")),
			mcp.WithNumber("limit", mcp.Description("返回的最大 key 数，默认1000，最大10000")),
		),
		handleRedisKeys,
	)

	s.AddTool(
		mcp.NewTool("data_redis_stream_range",
			mcp.WithDescription("读取 Stream 指定范围内的消息（只读）。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Stream key")),
			mcp.WithString("start", mcp.Description("起始 ID，默认 '-'（最早）")),
			mcp.WithString("stop", mcp.Description("结束 ID，默认 '+'（最新）")),
			mcp.WithNumber("count", mcp.Description("返回的最大消息数，默认100，最大1000")),
		),
		handleRedisStreamRange,
	)

	s.AddTool(
		mcp.NewTool("data_redis_stream_info",
			mcp.WithDescription("查看 Stream 元信息与消费者组信息（只读）：长度、首末消息、消费者组列表。"),
			mcp.WithString("key", mcp.Required(), mcp.Description("Stream key")),
		),
		handleRedisStreamInfo,
	)
}

// clampLimit 将 limit 限制在 [1, max]，非法值回退到 def。
func clampLimit(v, def, max int) int {
	if v < 1 {
		return def
	}
	if v > max {
		return max
	}
	return v
}

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", "auto")
	limit := clampLimit(int(req.GetFloat("limit", 100)), 100, 1000)

	// auto：先探测类型
	if keyType == "auto" {
		t, terr := redisx.Type(ctx, key)
		if terr != nil {
			return helper.ServiceError("探测Redis key类型", terr), nil
		}
		if t == "none" {
			return helper.EmptyResult(fmt.Sprintf("key 不存在: %s", key)), nil
		}
		keyType = t
	}

	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", "count": len(val), "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

	case "zset":
		val, err := redisx.ZRangeWithScores(ctx, key, 0, int64(limit-1))
		if err != nil {
			return helper.ServiceError("读取Redis zset", err), nil
		}
		if len(val) == 0 {
			return helper.EmptyResult(fmt.Sprintf("key 不存在或为空: %s", key)), nil
		}
		members := make([]map[string]any, 0, len(val))
		for _, z := range val {
			members = append(members, map[string]any{"member": z.Member, "score": z.Score})
		}
		return helper.JSONResultOrFileAuto(map[string]any{
			"key": key, "type": "zset", "count": len(members),
			"limit": limit, "truncated": len(members) == limit, "value": members,
		}), nil

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

func handleRedisMGet(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	raw, err := req.RequireString("keys")
	if err != nil {
		return helper.ParamError("keys"), nil
	}
	keys := make([]string, 0)
	for _, k := range strings.Split(raw, ",") {
		if s := strings.TrimSpace(k); s != "" {
			keys = append(keys, s)
		}
	}
	if len(keys) == 0 {
		return helper.ParamError("keys"), nil
	}
	vals, err := redisx.MGet(ctx, keys...)
	if err != nil {
		return helper.ServiceError("批量读取Redis string", err), nil
	}
	result := make(map[string]any, len(keys))
	for i, k := range keys {
		if i < len(vals) {
			result[k] = vals[i] // nil 表示 key 不存在
		}
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"type": "mget", "count": len(keys), "value": result,
	}), nil
}

func handleRedisHGet(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	key, err := req.RequireString("key")
	if err != nil {
		return helper.ParamError("key"), nil
	}
	field, err := req.RequireString("field")
	if err != nil {
		return helper.ParamError("field"), nil
	}
	val, err := redisx.HGet(ctx, key, field)
	if err != nil {
		return helper.EmptyResult(fmt.Sprintf("field 不存在或读取失败: %s.%s (%v)", key, field, err)), nil
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"key": key, "type": "hash", "field": field, "value": val,
	}), nil
}

func handleRedisType(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	key, err := req.RequireString("key")
	if err != nil {
		return helper.ParamError("key"), nil
	}
	t, err := redisx.Type(ctx, key)
	if err != nil {
		return helper.ServiceError("获取Redis key类型", err), nil
	}
	if t == "none" {
		return helper.EmptyResult(fmt.Sprintf("key 不存在: %s", key)), nil
	}

	out := map[string]any{"key": key, "type": t}

	// TTL
	if ttl, e := redisx.TTL(ctx, key); e == nil {
		out["ttl"] = ttlString(ttl.Seconds(), ttl == -1, ttl == -2)
	}

	// 各类型规模概览
	switch t {
	case "string":
		if n, e := redisx.StrLen(ctx, key); e == nil {
			out["length"] = n
		}
	case "hash":
		if n, e := redisx.HLen(ctx, key); e == nil {
			out["fields"] = n
		}
	case "list":
		if n, e := redisx.LLen(ctx, key); e == nil {
			out["length"] = n
		}
	case "set":
		if n, e := redisx.SCard(ctx, key); e == nil {
			out["members"] = n
		}
	case "zset":
		if n, e := redisx.ZCard(ctx, key); e == nil {
			out["members"] = n
		}
	case "stream":
		if n, e := redisx.XLen(ctx, key); e == nil {
			out["entries"] = n
		}
	}
	return helper.JSONResultOrFileAuto(out), 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
	}

	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
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"key": key, "exists": true,
		"ttl": ttlString(ttl.Seconds(), ttl == -1, ttl == -2),
	}), nil
}

func handleRedisTTL(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	key, err := req.RequireString("key")
	if err != nil {
		return helper.ParamError("key"), nil
	}
	ttl, err := redisx.TTL(ctx, key)
	if err != nil {
		return helper.ServiceError("获取TTL", err), nil
	}
	if ttl == -2 {
		return helper.EmptyResult(fmt.Sprintf("key 不存在: %s", key)), nil
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"key":         key,
		"ttl":         ttlString(ttl.Seconds(), ttl == -1, ttl == -2),
		"ttl_seconds": int64(ttl.Seconds()),
	}), nil
}

func handleRedisScan(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	cursor := uint64(req.GetFloat("cursor", 0))
	match := req.GetString("match", "")
	count := int64(clampLimit(int(req.GetFloat("count", 100)), 100, 1000))

	keys, next, err := redisx.Scan(ctx, cursor, match, count)
	if err != nil {
		return helper.ServiceError("SCAN 遍历", err), nil
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"type": "scan", "match": match, "count": len(keys),
		"next_cursor": next, "done": next == 0, "keys": keys,
	}), nil
}

func handleRedisKeys(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	pattern, err := req.RequireString("pattern")
	if err != nil {
		return helper.ParamError("pattern"), nil
	}
	limit := clampLimit(int(req.GetFloat("limit", 1000)), 1000, 10000)

	keys, err := redisx.Keys(ctx, pattern)
	if err != nil {
		return helper.ServiceError("KEYS 匹配", err), nil
	}
	total := len(keys)
	truncated := false
	if total > limit {
		keys = keys[:limit]
		truncated = true
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"type": "keys", "pattern": pattern, "total": total,
		"count": len(keys), "limit": limit, "truncated": truncated, "keys": keys,
	}), nil
}

func handleRedisStreamRange(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	key, err := req.RequireString("key")
	if err != nil {
		return helper.ParamError("key"), nil
	}
	start := req.GetString("start", "-")
	stop := req.GetString("stop", "+")
	count := int64(clampLimit(int(req.GetFloat("count", 100)), 100, 1000))

	msgs, err := redisx.XRangeN(ctx, key, start, stop, count)
	if err != nil {
		return helper.ServiceError("读取Stream消息", err), nil
	}
	if len(msgs) == 0 {
		return helper.EmptyResult(fmt.Sprintf("Stream 无消息或不存在: %s", key)), nil
	}
	items := make([]map[string]any, 0, len(msgs))
	for _, m := range msgs {
		items = append(items, map[string]any{"id": m.ID, "values": m.Values})
	}
	return helper.JSONResultOrFileAuto(map[string]any{
		"key": key, "type": "stream", "count": len(items),
		"limit": count, "truncated": int64(len(items)) == count, "value": items,
	}), nil
}

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

	xlen, err := redisx.XLen(ctx, key)
	if err != nil {
		return helper.ServiceError("获取Stream长度", err), nil
	}

	out := map[string]any{"key": key, "type": "stream", "length": xlen}

	if groups, e := redisx.XInfoGroups(ctx, key); e == nil {
		gs := make([]map[string]any, 0, len(groups))
		for _, g := range groups {
			gs = append(gs, map[string]any{
				"name":              g.Name,
				"consumers":         g.Consumers,
				"pending":           g.Pending,
				"last_delivered_id": g.LastDeliveredID,
			})
		}
		out["groups"] = gs
	}

	return helper.JSONResultOrFileAuto(out), nil
}

// ttlString 将 TTL 秒数格式化为可读字符串。
// persist 表示无过期时间(-1)，missing 表示 key 不存在(-2)。
func ttlString(seconds float64, persist, missing bool) string {
	if missing {
		return "key 不存在"
	}
	if persist || seconds < 0 {
		return "永不过期"
	}
	return fmt.Sprintf("%ds", int64(seconds))
}