sundynix-agentix/sundynix-gateway/internal/store/redis.go

package store

import (
	"context"
	"log"
	"time"

	"github.com/redis/go-redis/v9"
)

// Redis 持有 CacheDB 连接（Session / Rate Limit）。
// rdb 为 nil 表示降级模式（连接失败时放行，不限流）。
type Redis struct {
	rdb *redis.Client
}

// OpenRedis 连接 CacheDB。连接失败不 fatal：返回降级实例（限流放行）。
func OpenRedis(addr string) *Redis {
	rdb := redis.NewClient(&redis.Options{Addr: addr})
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := rdb.Ping(ctx).Err(); err != nil {
		log.Printf("[store] redis 不可用，降级运行（不限流）: %v", err)
		_ = rdb.Close()
		return &Redis{}
	}
	log.Println("[store] redis connected")
	return &Redis{rdb: rdb}
}

// Enabled 报告是否处于真实限流模式。
func (r *Redis) Enabled() bool { return r.rdb != nil }

// Allow 滑动窗口计数限流：在 window 内对 key 累加，超过 limit 即拒绝。
// 降级模式（rdb==nil）始终放行。
func (r *Redis) Allow(ctx context.Context, key string, limit int64, window time.Duration) (bool, error) {
	if r.rdb == nil {
		return true, nil
	}
	rk := "sundynix:ratelimit:" + key
	n, err := r.rdb.Incr(ctx, rk).Result()
	if err != nil {
		return true, err // 限流后端故障时放行，不阻断业务
	}
	if n == 1 {
		_ = r.rdb.Expire(ctx, rk, window).Err() // 首次计数设置窗口过期
	}
	return n <= limit, nil
}

// ---- Token 流持久化（Redis Stream：可回放的追加日志，根治 SSE 连晚/重连丢 token）----

// ---- Token 用量日预算（成本护栏：按用户按天累计 token，供提交前门控）----

// usageDailyKey 是某用户某天的累计 token key（day 形如 20260626，便于到期自然滚动）。
func usageDailyKey(userID, day string) string { return "sundynix:usage:" + userID + ":" + day }

// AddUsage 给某用户当天累计 token，并返回累计后总量（首次设 48h 过期自动清理）。降级返回 (0,nil)。
func (r *Redis) AddUsage(ctx context.Context, userID, day string, tokens int) (int64, error) {
	if r.rdb == nil || userID == "" {
		return 0, nil
	}
	rk := usageDailyKey(userID, day)
	n, err := r.rdb.IncrBy(ctx, rk, int64(tokens)).Result()
	if err != nil {
		return 0, err
	}
	if n == int64(tokens) { // 首次累计：设过期
		_ = r.rdb.Expire(ctx, rk, 48*time.Hour).Err()
	}
	return n, nil
}

// GetUsage 取某用户当天已用 token（不存在/降级返回 0）。
func (r *Redis) GetUsage(ctx context.Context, userID, day string) int64 {
	if r.rdb == nil || userID == "" {
		return 0
	}
	n, err := r.rdb.Get(ctx, usageDailyKey(userID, day)).Int64()
	if err != nil {
		return 0
	}
	return n
}

// ---- Token 流持久化（Redis Stream：可回放的追加日志，根治 SSE 连晚/重连丢 token）----

const streamTTL = 10 * time.Minute

func streamKey(taskID string) string { return "sundynix:stream:" + taskID }

// StreamEntry 是 token 流里的一条记录（ID 用于 SSE 的 Last-Event-ID 断点续传）。
type StreamEntry struct {
	ID   string
	Kind string // token / done
	Data string
}

// StreamAppend 把一条 token / 结束标记追加到任务的 Redis Stream（带 TTL 自动清理）。
func (r *Redis) StreamAppend(ctx context.Context, taskID, kind, data string) error {
	if r.rdb == nil {
		return nil
	}
	k := streamKey(taskID)
	if err := r.rdb.XAdd(ctx, &redis.XAddArgs{Stream: k, Values: map[string]any{"kind": kind, "data": data}}).Err(); err != nil {
		return err
	}
	return r.rdb.Expire(ctx, k, streamTTL).Err()
}

// StreamRead 从 lastID 之后阻塞读取新条目（XREAD BLOCK）。lastID="0" 表示从头回放。
// 阻塞超时无新数据时返回空切片 + 原 lastID（调用方据此继续轮询）。
func (r *Redis) StreamRead(ctx context.Context, taskID, lastID string, block time.Duration) ([]StreamEntry, string, error) {
	if r.rdb == nil {
		return nil, lastID, nil
	}
	res, err := r.rdb.XRead(ctx, &redis.XReadArgs{
		Streams: []string{streamKey(taskID), lastID}, Block: block, Count: 256,
	}).Result()
	if err == redis.Nil {
		return nil, lastID, nil // 阻塞超时、无新条目
	}
	if err != nil {
		return nil, lastID, err
	}
	var out []StreamEntry
	nl := lastID
	for _, st := range res {
		for _, m := range st.Messages {
			out = append(out, StreamEntry{ID: m.ID, Kind: asString(m.Values["kind"]), Data: asString(m.Values["data"])})
			nl = m.ID
		}
	}
	return out, nl, nil
}

func asString(v any) string {
	if s, ok := v.(string); ok {
		return s
	}
	return ""
}

// Close 释放底层连接。
func (r *Redis) Close() {
	if r.rdb != nil {
		_ = r.rdb.Close()
	}
}