internal/runtime/executor/cache_helpers.go

package executor

import (
	"context"
	json "github.com/goccy/go-json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
)

type codexCache struct {
	ID     string
	Expire time.Time
}

// codexCacheMap stores prompt cache IDs keyed by model+user_id.
// Protected by codexCacheMu. Entries expire after 1 hour.
var (
	codexCacheMap = make(map[string]codexCache)
	codexCacheMu  sync.RWMutex
	codexKVStore  cache.KVStore
)

// SetCodexKVStore configures an optional external store for codex prompt cache.
func SetCodexKVStore(store cache.KVStore) {
	codexKVStore = store
}

// codexCacheCleanupInterval controls how often expired entries are purged.
const codexCacheCleanupInterval = 15 * time.Minute

// codexCacheCleanupOnce ensures the background cleanup goroutine starts only once.
var codexCacheCleanupOnce sync.Once

// startCodexCacheCleanup launches a background goroutine that periodically
// removes expired entries from codexCacheMap to prevent memory leaks.
func startCodexCacheCleanup() {
	go func() {
		ticker := time.NewTicker(codexCacheCleanupInterval)
		defer ticker.Stop()
		for range ticker.C {
			purgeExpiredCodexCache()
		}
	}()
}

// purgeExpiredCodexCache removes entries that have expired.
func purgeExpiredCodexCache() {
	now := time.Now()
	codexCacheMu.Lock()
	defer codexCacheMu.Unlock()
	for key, cache := range codexCacheMap {
		if cache.Expire.Before(now) {
			delete(codexCacheMap, key)
		}
	}
}

// getCodexCache retrieves a cached entry, returning ok=false if not found or expired.
func getCodexCache(key string) (codexCache, bool) {
	if codexKVStore != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		payload, err := codexKVStore.Get(ctx, codexRedisKey(key))
		cancel()
		if err == nil && payload != "" {
			var entry codexCache
			if json.Unmarshal([]byte(payload), &entry) == nil && !entry.Expire.Before(time.Now()) {
				return entry, true
			}
		}
	}

	codexCacheCleanupOnce.Do(startCodexCacheCleanup)
	codexCacheMu.RLock()
	cache, ok := codexCacheMap[key]
	codexCacheMu.RUnlock()
	if !ok || cache.Expire.Before(time.Now()) {
		return codexCache{}, false
	}
	return cache, true
}

// setCodexCache stores a cache entry.
func setCodexCache(key string, cache codexCache) {
	if codexKVStore != nil {
		if payload, err := json.Marshal(cache); err == nil {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			ttl := time.Until(cache.Expire)
			if ttl <= 0 {
				ttl = time.Hour
			}
			_ = codexKVStore.Set(ctx, codexRedisKey(key), string(payload), ttl)
			cancel()
		}
	}

	codexCacheCleanupOnce.Do(startCodexCacheCleanup)
	codexCacheMu.Lock()
	codexCacheMap[key] = cache
	codexCacheMu.Unlock()
}

func codexRedisKey(key string) string {
	trimmed := strings.TrimSpace(key)
	if trimmed == "" {
		return "codex"
	}
	parts := strings.SplitN(trimmed, ":", 2)
	if len(parts) == 2 && strings.TrimSpace(parts[0]) != "" && strings.TrimSpace(parts[1]) != "" {
		return fmt.Sprintf("codex:%s:%s", strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1]))
	}
	return fmt.Sprintf("codex:%s", trimmed)
}