Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
3.8 kB · 148 lines
Go
at commit e8e39526
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149package helps
import ( "context" "crypto/sha256" "encoding/hex" "fmt" "strings" "sync" "time"
"github.com/google/uuid" homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home")
type sessionIDCacheEntry struct { value string expire time.Time}
var ( sessionIDCache = make(map[string]sessionIDCacheEntry) sessionIDCacheMu sync.RWMutex sessionIDCacheCleanupOnce sync.Once)
type claudeIDKVClient interface { KVGet(ctx context.Context, key string) ([]byte, bool, error) KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error)}
var currentClaudeIDKVClient = func() (claudeIDKVClient, bool, error) { return homekv.CurrentKVClient()}
const ( sessionIDTTL = time.Hour sessionIDCacheCleanupPeriod = 15 * time.Minute)
func startSessionIDCacheCleanup() { go func() { ticker := time.NewTicker(sessionIDCacheCleanupPeriod) defer ticker.Stop() for range ticker.C { purgeExpiredSessionIDs() } }()}
func purgeExpiredSessionIDs() { now := time.Now() sessionIDCacheMu.Lock() for key, entry := range sessionIDCache { if !entry.expire.After(now) { delete(sessionIDCache, key) } } sessionIDCacheMu.Unlock()}
func sessionIDCacheKey(apiKey string) string { sum := sha256.Sum256([]byte(apiKey)) return hex.EncodeToString(sum[:])}
// CachedSessionID returns a stable session UUID per apiKey, refreshing the TTL on each access.func CachedSessionID(apiKey string) string { value, errValue := CachedSessionIDRequired(context.Background(), apiKey) if errValue == nil && value != "" { return value } return uuid.New().String()}
// CachedSessionIDRequired returns a stable session UUID per apiKey for request-time paths.func CachedSessionIDRequired(ctx context.Context, apiKey string) (string, error) { if apiKey == "" { return uuid.New().String(), nil } client, homeMode, errClient := currentClaudeIDKVClient() if homeMode { if errClient != nil { return "", errClient } key := claudeSessionIDKVKey(apiKey) raw, found, errGet := client.KVGet(ctx, key) if errGet != nil { return "", errGet } if found && strings.TrimSpace(string(raw)) != "" { if _, errExpire := client.KVExpire(ctx, key, sessionIDTTL); errExpire != nil { return "", errExpire } return strings.TrimSpace(string(raw)), nil } newID := uuid.New().String() if _, errSet := client.KVSetNX(ctx, key, []byte(newID), sessionIDTTL); errSet != nil { return "", errSet } raw, found, errGet = client.KVGet(ctx, key) if errGet != nil { return "", errGet } if found && strings.TrimSpace(string(raw)) != "" { return strings.TrimSpace(string(raw)), nil } return "", fmt.Errorf("home kv session id missing after set") }
sessionIDCacheCleanupOnce.Do(startSessionIDCacheCleanup)
key := sessionIDCacheKey(apiKey) now := time.Now()
sessionIDCacheMu.RLock() entry, ok := sessionIDCache[key] valid := ok && entry.value != "" && entry.expire.After(now) sessionIDCacheMu.RUnlock() if valid { sessionIDCacheMu.Lock() entry = sessionIDCache[key] if entry.value != "" && entry.expire.After(now) { entry.expire = now.Add(sessionIDTTL) sessionIDCache[key] = entry sessionIDCacheMu.Unlock() return entry.value, nil } sessionIDCacheMu.Unlock() }
newID := uuid.New().String()
sessionIDCacheMu.Lock() entry, ok = sessionIDCache[key] if !ok || entry.value == "" || !entry.expire.After(now) { entry.value = newID } entry.expire = now.Add(sessionIDTTL) sessionIDCache[key] = entry sessionIDCacheMu.Unlock() return entry.value, nil}
func claudeSessionIDKVKey(apiKey string) string { return "cpa:claude:session-id:" + homekv.HashKeyPart(apiKey)}