internal/cache/response_cache.go

package cache

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"path/filepath"
	"sync/atomic"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
)

// CachedResponse is a stored API response.
type CachedResponse struct {
	StatusCode int               `json:"s"`
	Headers    map[string]string `json:"h"`
	Body       []byte            `json:"b"`
	CachedAt   time.Time         `json:"t"`
	Model      string            `json:"m"`
}

// CacheStats holds cache hit/miss counters.
type CacheStats struct {
	Hits   int64 `json:"hits"`
	Misses int64 `json:"misses"`
	Size   int64 `json:"size"`
}

// ResponseCache defines the response caching interface.
type ResponseCache interface {
	Get(ctx context.Context, key string) (*CachedResponse, bool)
	Set(ctx context.Context, key string, resp *CachedResponse, ttl time.Duration) error
	Delete(ctx context.Context, pattern string) error
	Stats() CacheStats
	Close() error
}

// ResponseCacheKey generates a cache key from model name and request body.
func ResponseCacheKey(model string, body []byte) string {
	h := sha256.New()
	h.Write([]byte(model))
	h.Write([]byte{0})
	h.Write(body)
	return hex.EncodeToString(h.Sum(nil))
}

// ShouldCache returns whether a response should be cached based on config rules.
func ShouldCache(cfg config.ResponseCacheConfig, model string, isStreaming bool, responseSize int) bool {
	if !cfg.Enabled {
		return false
	}
	if isStreaming && cfg.ExcludeStreaming {
		return false
	}
	if cfg.MaxResponseSizeKB > 0 && responseSize > cfg.MaxResponseSizeKB*1024 {
		return false
	}
	for _, pattern := range cfg.ExcludeModels {
		if matched, _ := filepath.Match(pattern, model); matched {
			return false
		}
	}
	return true
}

// NewResponseCache creates a ResponseCache backed by Redis (if client != nil) or in-memory.
func NewResponseCache(cfg config.ResponseCacheConfig, client redis.UniversalClient) ResponseCache {
	if client != nil {
		return newRedisResponseCache(cfg, client)
	}
	return newMemoryResponseCache(cfg)
}

// responseCacheStats provides thread-safe hit/miss counting.
type responseCacheStats struct {
	hits   atomic.Int64
	misses atomic.Int64
}

func (s *responseCacheStats) recordHit()  { s.hits.Add(1) }
func (s *responseCacheStats) recordMiss() { s.misses.Add(1) }