From 2a050dc95d418f15555d49b1c235deb181a07433 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sun, 14 Jun 2026 20:32:47 +0800 Subject: [PATCH] feat: enhance fault tolerance for kv-based caching and introduce additional tests - Updated Antigravity Credits fallback to handle KV store unavailability as a service error. - Enhanced signature caching mechanisms with request-time KV access and sliding expiration. - Added and improved tests for KV client interactions, including error handling and expiration behaviors. - Introduced `CacheSignatureBestEffort` for non-critical signature caching and clarified function flows with required context. - Ensured consistent error reporting for missing or unavailable KV stores in various scenarios. - Replaced direct `homekv` calls with injectable KV client interfaces for `antigravity` and `codex_reasoning_replay` modules. - Improved error reporting and handling for KV operations, including `KVGet`, `KVSet`, `KVDel`, and `KVExpire`. - Introduced dedicated fake KV clients for expanded and granular test coverage. - Added new unit tests to validate KV client behaviors and error scenarios, ensuring robustness and sliding expiration functionality. --- .../cache/codex_reasoning_replay_cache.go | 94 ++++++- .../codex_reasoning_replay_cache_test.go | 176 ++++++++++++ internal/cache/signature_cache.go | 119 +++++++- internal/cache/signature_cache_test.go | 200 +++++++++++++ internal/home/client.go | 188 +++++++++++++ internal/home/client_test.go | 243 ++++++++++++++++ internal/home/kv_helpers.go | 189 +++++++++++++ internal/home/kv_helpers_test.go | 110 ++++++++ .../runtime/executor/antigravity_executor.go | 264 ++++++++++++++++-- .../antigravity_executor_credits_test.go | 230 +++++++++++++++ .../antigravity_executor_signature_test.go | 8 +- internal/runtime/executor/claude_executor.go | 76 +++-- .../runtime/executor/claude_executor_test.go | 5 +- internal/runtime/executor/codex_executor.go | 82 ++++-- .../executor/codex_websockets_executor.go | 25 +- .../runtime/executor/helps/cache_helpers.go | 64 ++++- .../executor/helps/cache_helpers_test.go | 27 ++ .../executor/helps/claude_device_profile.go | 169 +++++++++++ .../helps/claude_device_profile_test.go | 237 ++++++++++++++++ .../executor/helps/session_id_cache.go | 62 +++- .../executor/helps/session_id_cache_test.go | 178 ++++++++++++ .../runtime/executor/helps/user_id_cache.go | 53 +++- .../executor/helps/user_id_cache_test.go | 79 ++++++ .../claude/antigravity_claude_request.go | 79 +++++- .../claude/antigravity_claude_response.go | 4 +- sdk/cliproxy/auth/antigravity_credits.go | 32 ++- sdk/cliproxy/auth/antigravity_credits_test.go | 30 ++ sdk/cliproxy/auth/conductor.go | 84 +++--- .../auth/conductor_credits_candidates_test.go | 45 ++- 29 files changed, 2991 insertions(+), 161 deletions(-) create mode 100644 internal/home/kv_helpers.go create mode 100644 internal/home/kv_helpers_test.go create mode 100644 internal/runtime/executor/helps/cache_helpers_test.go create mode 100644 internal/runtime/executor/helps/claude_device_profile_test.go create mode 100644 internal/runtime/executor/helps/session_id_cache_test.go diff --git a/internal/cache/codex_reasoning_replay_cache.go b/internal/cache/codex_reasoning_replay_cache.go index 820f7f1d..274d131b 100644 --- a/internal/cache/codex_reasoning_replay_cache.go +++ b/internal/cache/codex_reasoning_replay_cache.go @@ -1,12 +1,16 @@ package cache import ( + "context" + "encoding/json" "sort" "strings" "sync" "time" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" + log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -35,6 +39,17 @@ var ( codexReasoningReplayEntries = make(map[string]codexReasoningReplayEntry) ) +type codexReasoningReplayKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVDel(ctx context.Context, keys ...string) (int64, error) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentCodexReasoningReplayKVClient = func() (codexReasoningReplayKVClient, bool, error) { + return homekv.CurrentKVClient() +} + // CacheCodexReasoningReplayItem stores a final GPT/Codex reasoning item for // stateless replay. The stored item is normalized to the minimal shape accepted // by Responses input replay. @@ -45,6 +60,11 @@ func CacheCodexReasoningReplayItem(modelName, sessionKey string, item []byte) bo // CacheCodexReasoningReplayItems stores the final GPT/Codex assistant output // items needed to replay a stateless next turn. func CacheCodexReasoningReplayItems(modelName, sessionKey string, items [][]byte) bool { + return CacheCodexReasoningReplayItemsBestEffort(context.Background(), modelName, sessionKey, items) +} + +// CacheCodexReasoningReplayItemsBestEffort stores replay items for completed response paths. +func CacheCodexReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool { key := codexReasoningReplayCacheKey(modelName, sessionKey) if key == "" { return false @@ -53,6 +73,23 @@ func CacheCodexReasoningReplayItems(modelName, sessionKey string, items [][]byte if !ok { return false } + if client, homeMode, errClient := currentCodexReasoningReplayKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errClient) + return false + } + raw, errMarshal := json.Marshal(normalized) + if errMarshal != nil { + log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errMarshal) + return false + } + written, errSet := client.KVSet(ctx, codexReasoningReplayKVKey(modelName, sessionKey), raw, homekv.KVSetOptions{EX: CodexReasoningReplayCacheTTL}) + if errSet != nil { + log.Errorf("home kv best-effort codex reasoning replay set failed prefix=cpa:codex:*: %v", errSet) + return false + } + return written + } cacheCleanupOnce.Do(startCacheCleanup) now := time.Now() @@ -79,9 +116,36 @@ func GetCodexReasoningReplayItem(modelName, sessionKey string) ([]byte, bool) { // GetCodexReasoningReplayItems retrieves normalized assistant output items. func GetCodexReasoningReplayItems(modelName, sessionKey string) ([][]byte, bool) { + items, ok, err := GetCodexReasoningReplayItemsRequired(context.Background(), modelName, sessionKey) + if err == nil { + return items, ok + } + return nil, false +} + +// GetCodexReasoningReplayItemsRequired retrieves replay items for request-time paths. +func GetCodexReasoningReplayItemsRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, bool, error) { key := codexReasoningReplayCacheKey(modelName, sessionKey) if key == "" { - return nil, false + return nil, false, nil + } + client, homeMode, errClient := currentCodexReasoningReplayKVClient() + if homeMode { + if errClient != nil { + return nil, false, errClient + } + raw, found, errGet := client.KVGet(ctx, codexReasoningReplayKVKey(modelName, sessionKey)) + if errGet != nil || !found { + return nil, false, errGet + } + var homeItems [][]byte + if errUnmarshal := json.Unmarshal(raw, &homeItems); errUnmarshal != nil { + return nil, false, errUnmarshal + } + if _, errExpire := client.KVExpire(ctx, codexReasoningReplayKVKey(modelName, sessionKey), CodexReasoningReplayCacheTTL); errExpire != nil { + return nil, false, errExpire + } + return cloneCodexReasoningReplayItems(homeItems), true, nil } cacheCleanupOnce.Do(startCacheCleanup) @@ -90,27 +154,43 @@ func GetCodexReasoningReplayItems(modelName, sessionKey string) ([][]byte, bool) defer codexReasoningReplayMu.Unlock() entry, ok := codexReasoningReplayEntries[key] if !ok { - return nil, false + return nil, false, nil } if now.Sub(entry.Timestamp) > CodexReasoningReplayCacheTTL { delete(codexReasoningReplayEntries, key) - return nil, false + return nil, false, nil } entry.Timestamp = now codexReasoningReplayEntries[key] = entry - return cloneCodexReasoningReplayItems(entry.Items), true + return cloneCodexReasoningReplayItems(entry.Items), true, nil } // DeleteCodexReasoningReplayItem removes one replay item after upstream rejects // it or the caller otherwise knows it is stale. func DeleteCodexReasoningReplayItem(modelName, sessionKey string) { + if errDelete := DeleteCodexReasoningReplayItemRequired(context.Background(), modelName, sessionKey); errDelete != nil { + return + } +} + +// DeleteCodexReasoningReplayItemRequired removes one replay item for request-time paths. +func DeleteCodexReasoningReplayItemRequired(ctx context.Context, modelName, sessionKey string) error { key := codexReasoningReplayCacheKey(modelName, sessionKey) if key == "" { - return + return nil + } + client, homeMode, errClient := currentCodexReasoningReplayKVClient() + if homeMode { + if errClient != nil { + return errClient + } + _, errDel := client.KVDel(ctx, codexReasoningReplayKVKey(modelName, sessionKey)) + return errDel } codexReasoningReplayMu.Lock() delete(codexReasoningReplayEntries, key) codexReasoningReplayMu.Unlock() + return nil } // ClearCodexReasoningReplayCache clears all Codex reasoning replay state. @@ -131,6 +211,10 @@ func codexReasoningReplayCacheKey(modelName, sessionKey string) string { return strings.Join([]string{"codex-reasoning-replay", modelName, sessionKey}, "\x00") } +func codexReasoningReplayKVKey(modelName, sessionKey string) string { + return "cpa:codex:reasoning-replay:" + homekv.HashKeyPart(strings.TrimSpace(modelName)) + ":" + homekv.HashKeyPart(strings.TrimSpace(sessionKey)) +} + func normalizeCodexReasoningReplayItems(items [][]byte) ([][]byte, bool) { normalized := make([][]byte, 0, len(items)) for _, item := range items { diff --git a/internal/cache/codex_reasoning_replay_cache_test.go b/internal/cache/codex_reasoning_replay_cache_test.go index cc43ed41..8bfe494f 100644 --- a/internal/cache/codex_reasoning_replay_cache_test.go +++ b/internal/cache/codex_reasoning_replay_cache_test.go @@ -1,11 +1,92 @@ package cache import ( + "context" "encoding/base64" + "encoding/json" + "errors" "fmt" "testing" + "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" ) +type fakeCodexReasoningReplayKVClient struct { + values map[string][]byte + getErr error + setErr error + delErr error + expireErr error + getCount int + setCount int + delCount int + expireCount int + lastSetTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeCodexReasoningReplayKVClient() *fakeCodexReasoningReplayKVClient { + return &fakeCodexReasoningReplayKVClient{values: make(map[string][]byte)} +} + +func (c *fakeCodexReasoningReplayKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeCodexReasoningReplayKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.setCount++ + c.lastSetTTL = opts.EX + if c.setErr != nil { + return false, c.setErr + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + +func (c *fakeCodexReasoningReplayKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { + c.delCount++ + if c.delErr != nil { + return 0, c.delErr + } + var deleted int64 + for _, key := range keys { + if _, ok := c.values[key]; ok { + delete(c.values, key) + deleted++ + } + } + return deleted, nil +} + +func (c *fakeCodexReasoningReplayKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.expireCount++ + c.lastExpireTTL = ttl + if c.expireErr != nil { + return false, c.expireErr + } + return true, nil +} + +func useFakeCodexReasoningReplayKVClient(t *testing.T, client *fakeCodexReasoningReplayKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentCodexReasoningReplayKVClient + currentCodexReasoningReplayKVClient = func() (codexReasoningReplayKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentCodexReasoningReplayKVClient = previous + }) +} + func validCodexReasoningReplayEncryptedContentForTest(seed byte) string { payload := make([]byte, 1+8+16+16+32) payload[0] = 0x80 @@ -15,6 +96,19 @@ func validCodexReasoningReplayEncryptedContentForTest(seed byte) string { return base64.RawURLEncoding.EncodeToString(payload) } +func validCodexReasoningReplayItemForTest(seed byte) []byte { + return []byte(`{"type":"reasoning","summary":[],"content":null,"encrypted_content":"` + validCodexReasoningReplayEncryptedContentForTest(seed) + `"}`) +} + +func mustCodexReasoningReplayJSON(t *testing.T, items [][]byte) []byte { + t.Helper() + raw, errMarshal := json.Marshal(items) + if errMarshal != nil { + t.Fatalf("marshal replay items: %v", errMarshal) + } + return raw +} + func TestCodexReasoningReplayCacheRejectsInvalidItems(t *testing.T) { ClearCodexReasoningReplayCache() t.Cleanup(ClearCodexReasoningReplayCache) @@ -27,6 +121,88 @@ func TestCodexReasoningReplayCacheRejectsInvalidItems(t *testing.T) { } } +func TestCodexReasoningReplayRequiredHomeReadAndSlidingExpire(t *testing.T) { + ClearCodexReasoningReplayCache() + t.Cleanup(ClearCodexReasoningReplayCache) + client := newFakeCodexReasoningReplayKVClient() + key := codexReasoningReplayKVKey("gpt-5.4", "session-home") + item := validCodexReasoningReplayItemForTest(3) + client.values[key] = mustCodexReasoningReplayJSON(t, [][]byte{item}) + useFakeCodexReasoningReplayKVClient(t, client, true, nil) + + items, found, errGet := GetCodexReasoningReplayItemsRequired(context.Background(), "gpt-5.4", "session-home") + if errGet != nil { + t.Fatalf("GetCodexReasoningReplayItemsRequired() error = %v", errGet) + } + if !found || len(items) != 1 || string(items[0]) != string(item) { + t.Fatalf("GetCodexReasoningReplayItemsRequired() = %q, %v, want item, true", items, found) + } + if client.expireCount != 1 || client.lastExpireTTL != CodexReasoningReplayCacheTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, CodexReasoningReplayCacheTTL) + } +} + +func TestCodexReasoningReplayRequiredHomeFailures(t *testing.T) { + for _, tc := range []struct { + name string + client *fakeCodexReasoningReplayKVClient + }{ + {name: "get", client: &fakeCodexReasoningReplayKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "expire", client: &fakeCodexReasoningReplayKVClient{values: map[string][]byte{ + codexReasoningReplayKVKey("gpt-5.4", "session-home"): mustCodexReasoningReplayJSON(t, [][]byte{validCodexReasoningReplayItemForTest(4)}), + }, expireErr: errors.New("expire failed")}}, + {name: "delete", client: &fakeCodexReasoningReplayKVClient{values: make(map[string][]byte), delErr: errors.New("delete failed")}}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeCodexReasoningReplayKVClient(t, tc.client, true, nil) + switch tc.name { + case "delete": + if errDel := DeleteCodexReasoningReplayItemRequired(context.Background(), "gpt-5.4", "session-home"); errDel == nil { + t.Fatalf("DeleteCodexReasoningReplayItemRequired() error = nil, want error") + } + default: + if _, _, errGet := GetCodexReasoningReplayItemsRequired(context.Background(), "gpt-5.4", "session-home"); errGet == nil { + t.Fatalf("GetCodexReasoningReplayItemsRequired() error = nil, want error") + } + } + }) + } +} + +func TestCodexReasoningReplayBestEffortHomeWriteFailureDoesNotUseLocalCache(t *testing.T) { + ClearCodexReasoningReplayCache() + t.Cleanup(ClearCodexReasoningReplayCache) + client := newFakeCodexReasoningReplayKVClient() + client.setErr = errors.New("set failed") + useFakeCodexReasoningReplayKVClient(t, client, true, nil) + + if CacheCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "session-home", [][]byte{validCodexReasoningReplayItemForTest(5)}) { + t.Fatalf("CacheCodexReasoningReplayItemsBestEffort() = true, want false") + } + useFakeCodexReasoningReplayKVClient(t, newFakeCodexReasoningReplayKVClient(), false, nil) + if _, found := GetCodexReasoningReplayItems("gpt-5.4", "session-home"); found { + t.Fatalf("local replay cache was populated after Home best-effort write failure") + } +} + +func TestCodexReasoningReplayHomeRejectsEmptyScopeWithoutKV(t *testing.T) { + client := newFakeCodexReasoningReplayKVClient() + useFakeCodexReasoningReplayKVClient(t, client, true, nil) + + if _, found, errGet := GetCodexReasoningReplayItemsRequired(context.Background(), "", "session-home"); errGet != nil || found { + t.Fatalf("GetCodexReasoningReplayItemsRequired(empty model) = found %v err %v, want false nil", found, errGet) + } + if CacheCodexReasoningReplayItemsBestEffort(context.Background(), "gpt-5.4", "", [][]byte{validCodexReasoningReplayItemForTest(6)}) { + t.Fatalf("CacheCodexReasoningReplayItemsBestEffort(empty session) = true, want false") + } + if errDel := DeleteCodexReasoningReplayItemRequired(context.Background(), "gpt-5.4", ""); errDel != nil { + t.Fatalf("DeleteCodexReasoningReplayItemRequired(empty session) error = %v", errDel) + } + if client.getCount != 0 || client.setCount != 0 || client.delCount != 0 || client.expireCount != 0 { + t.Fatalf("KV calls = get %d set %d del %d expire %d, want all zero", client.getCount, client.setCount, client.delCount, client.expireCount) + } +} + func TestCodexReasoningReplayCacheScopesByModelAndSession(t *testing.T) { ClearCodexReasoningReplayCache() t.Cleanup(ClearCodexReasoningReplayCache) diff --git a/internal/cache/signature_cache.go b/internal/cache/signature_cache.go index 42020ae7..1f54458e 100644 --- a/internal/cache/signature_cache.go +++ b/internal/cache/signature_cache.go @@ -1,13 +1,16 @@ package cache import ( + "context" "crypto/sha256" "encoding/hex" + "fmt" "strings" "sync" "sync/atomic" "time" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" log "github.com/sirupsen/logrus" ) @@ -37,6 +40,17 @@ var signatureCache sync.Map // cacheCleanupOnce ensures the background cleanup goroutine starts only once var cacheCleanupOnce sync.Once +type signatureKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVDel(ctx context.Context, keys ...string) (int64, error) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentSignatureKVClient = func() (signatureKVClient, bool, error) { + return homekv.CurrentKVClient() +} + // groupCache is the inner map type type groupCache struct { mu sync.RWMutex @@ -100,11 +114,29 @@ func purgeExpiredCaches() { // CacheSignature stores a thinking signature for a given model group and text. // Used for Claude models that require signed thinking blocks in multi-turn conversations. func CacheSignature(modelName, text, signature string) { + CacheSignatureBestEffort(context.Background(), modelName, text, signature) +} + +// CacheSignatureBestEffort stores a thinking signature for completed response paths. +func CacheSignatureBestEffort(ctx context.Context, modelName, text, signature string) bool { if text == "" || signature == "" { - return + return false } if len(signature) < MinValidSignatureLen { - return + return false + } + + if client, homeMode, errClient := currentSignatureKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort signature set failed prefix=cpa:signature:*: %v", errClient) + return false + } + written, errSet := client.KVSet(ctx, signatureKVKey(modelName, text), []byte(signature), homekv.KVSetOptions{EX: SignatureCacheTTL}) + if errSet != nil { + log.Errorf("home kv best-effort signature set failed prefix=cpa:signature:*: %v", errSet) + return false + } + return written } groupKey := GetModelGroup(modelName) @@ -117,25 +149,57 @@ func CacheSignature(modelName, text, signature string) { Signature: signature, Timestamp: time.Now(), } + return true } // GetCachedSignature retrieves a cached signature for a given model group and text. // Returns empty string if not found or expired. func GetCachedSignature(modelName, text string) string { + signature, errSignature := GetCachedSignatureRequired(context.Background(), modelName, text) + if errSignature != nil { + return "" + } + return signature +} + +// GetCachedSignatureRequired retrieves a cached signature for request-time paths. +func GetCachedSignatureRequired(ctx context.Context, modelName, text string) (string, error) { groupKey := GetModelGroup(modelName) if text == "" { if groupKey == "gemini" { - return "skip_thought_signature_validator" + return "skip_thought_signature_validator", nil } - return "" + return "", nil } + + if client, homeMode, errClient := currentSignatureKVClient(); homeMode { + if errClient != nil { + return "", errClient + } + key := signatureKVKey(modelName, text) + raw, found, errGet := client.KVGet(ctx, key) + if errGet != nil { + return "", errGet + } + if !found { + if groupKey == "gemini" { + return "skip_thought_signature_validator", nil + } + return "", nil + } + if _, errExpire := client.KVExpire(ctx, key, SignatureCacheTTL); errExpire != nil { + return "", errExpire + } + return string(raw), nil + } + val, ok := signatureCache.Load(groupKey) if !ok { if groupKey == "gemini" { - return "skip_thought_signature_validator" + return "skip_thought_signature_validator", nil } - return "" + return "", nil } sc := val.(*groupCache) @@ -148,17 +212,17 @@ func GetCachedSignature(modelName, text string) string { if !exists { sc.mu.Unlock() if groupKey == "gemini" { - return "skip_thought_signature_validator" + return "skip_thought_signature_validator", nil } - return "" + return "", nil } if now.Sub(entry.Timestamp) > SignatureCacheTTL { delete(sc.entries, textHash) sc.mu.Unlock() if groupKey == "gemini" { - return "skip_thought_signature_validator" + return "skip_thought_signature_validator", nil } - return "" + return "", nil } // Refresh TTL on access (sliding expiration). @@ -166,7 +230,7 @@ func GetCachedSignature(modelName, text string) string { sc.entries[textHash] = entry sc.mu.Unlock() - return entry.Signature + return entry.Signature, nil } // ClearSignatureCache clears signature cache for a specific model group or all groups. @@ -182,6 +246,35 @@ func ClearSignatureCache(modelName string) { signatureCache.Delete(groupKey) } +// DeleteCachedSignatureRequired removes one exact cached signature. +func DeleteCachedSignatureRequired(ctx context.Context, modelName, text string) error { + if text == "" { + return nil + } + if client, homeMode, errClient := currentSignatureKVClient(); homeMode { + if errClient != nil { + return errClient + } + _, errDel := client.KVDel(ctx, signatureKVKey(modelName, text)) + return errDel + } + groupKey := GetModelGroup(modelName) + textHash := hashText(text) + val, ok := signatureCache.Load(groupKey) + if !ok { + return nil + } + sc := val.(*groupCache) + sc.mu.Lock() + delete(sc.entries, textHash) + isEmpty := len(sc.entries) == 0 + sc.mu.Unlock() + if isEmpty { + signatureCache.Delete(groupKey) + } + return nil +} + // HasValidSignature checks if a signature is valid (non-empty and long enough) func HasValidSignature(modelName, signature string) bool { return (signature != "" && len(signature) >= MinValidSignatureLen) || (signature == "skip_thought_signature_validator" && GetModelGroup(modelName) == "gemini") @@ -198,6 +291,10 @@ func GetModelGroup(modelName string) string { return modelName } +func signatureKVKey(modelName, text string) string { + return fmt.Sprintf("cpa:signature:%s:%s", GetModelGroup(modelName), homekv.HashKeyPart(text)) +} + var signatureCacheEnabled atomic.Bool var signatureBypassStrictMode atomic.Bool diff --git a/internal/cache/signature_cache_test.go b/internal/cache/signature_cache_test.go index 82a8a19d..5fe5b9e0 100644 --- a/internal/cache/signature_cache_test.go +++ b/internal/cache/signature_cache_test.go @@ -2,15 +2,93 @@ package cache import ( "bytes" + "context" + "errors" "strings" "testing" "time" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" log "github.com/sirupsen/logrus" ) const testModelName = "claude-sonnet-4-5" +type fakeSignatureKVClient struct { + values map[string][]byte + getErr error + setErr error + delErr error + expireErr error + getCount int + setCount int + delCount int + expireCount int + lastSetTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeSignatureKVClient() *fakeSignatureKVClient { + return &fakeSignatureKVClient{values: make(map[string][]byte)} +} + +func (c *fakeSignatureKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeSignatureKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.setCount++ + c.lastSetTTL = opts.EX + if c.setErr != nil { + return false, c.setErr + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + +func (c *fakeSignatureKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { + c.delCount++ + if c.delErr != nil { + return 0, c.delErr + } + var deleted int64 + for _, key := range keys { + if _, ok := c.values[key]; ok { + delete(c.values, key) + deleted++ + } + } + return deleted, nil +} + +func (c *fakeSignatureKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.expireCount++ + c.lastExpireTTL = ttl + if c.expireErr != nil { + return false, c.expireErr + } + return true, nil +} + +func useFakeSignatureKVClient(t *testing.T, client *fakeSignatureKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentSignatureKVClient + currentSignatureKVClient = func() (signatureKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentSignatureKVClient = previous + }) +} + func TestCacheSignature_BasicStorageAndRetrieval(t *testing.T) { ClearSignatureCache("") @@ -27,6 +105,128 @@ func TestCacheSignature_BasicStorageAndRetrieval(t *testing.T) { } } +func TestGetCachedSignatureRequiredHomeReadAndSlidingExpire(t *testing.T) { + ClearSignatureCache("") + text := "thinking text" + signature := "abc123validSignature1234567890123456789012345678901234567890" + client := newFakeSignatureKVClient() + client.values[signatureKVKey(testModelName, text)] = []byte(signature) + useFakeSignatureKVClient(t, client, true, nil) + + got, errGet := GetCachedSignatureRequired(context.Background(), testModelName, text) + if errGet != nil { + t.Fatalf("GetCachedSignatureRequired() error = %v", errGet) + } + if got != signature { + t.Fatalf("GetCachedSignatureRequired() = %q, want %q", got, signature) + } + if client.expireCount != 1 || client.lastExpireTTL != SignatureCacheTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, SignatureCacheTTL) + } +} + +func TestGetCachedSignatureRequiredHomeFailures(t *testing.T) { + for _, tc := range []struct { + name string + client *fakeSignatureKVClient + }{ + {name: "get", client: &fakeSignatureKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "expire", client: &fakeSignatureKVClient{values: map[string][]byte{ + signatureKVKey(testModelName, "thinking text"): []byte("abc123validSignature1234567890123456789012345678901234567890"), + }, expireErr: errors.New("expire failed")}}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeSignatureKVClient(t, tc.client, true, nil) + if _, errGet := GetCachedSignatureRequired(context.Background(), testModelName, "thinking text"); errGet == nil { + t.Fatalf("GetCachedSignatureRequired() error = nil, want error") + } + }) + } +} + +func TestGetCachedSignatureRequiredHomeMissDoesNotFallbackToLocalCache(t *testing.T) { + ClearSignatureCache("") + text := "thinking text" + signature := "abc123validSignature1234567890123456789012345678901234567890" + CacheSignature(testModelName, text, signature) + + client := newFakeSignatureKVClient() + useFakeSignatureKVClient(t, client, true, nil) + + got, errGet := GetCachedSignatureRequired(context.Background(), testModelName, text) + if errGet != nil { + t.Fatalf("GetCachedSignatureRequired() error = %v", errGet) + } + if got != "" { + t.Fatalf("GetCachedSignatureRequired() = %q, want Home miss without local fallback", got) + } +} + +func TestCacheSignatureBestEffortHomeWriteFailureDoesNotUseLocalCache(t *testing.T) { + ClearSignatureCache("") + text := "thinking text" + signature := "abc123validSignature1234567890123456789012345678901234567890" + client := newFakeSignatureKVClient() + client.setErr = errors.New("set failed") + useFakeSignatureKVClient(t, client, true, nil) + + if CacheSignatureBestEffort(context.Background(), testModelName, text, signature) { + t.Fatalf("CacheSignatureBestEffort() = true, want false") + } + useFakeSignatureKVClient(t, newFakeSignatureKVClient(), false, nil) + if got := GetCachedSignature(testModelName, text); got != "" { + t.Fatalf("local cache = %q, want empty after Home write failure", got) + } +} + +func TestDeleteCachedSignatureRequiredHomeExactKey(t *testing.T) { + ClearSignatureCache("") + text := "thinking text" + signature := "abc123validSignature1234567890123456789012345678901234567890" + client := newFakeSignatureKVClient() + client.values[signatureKVKey(testModelName, text)] = []byte(signature) + useFakeSignatureKVClient(t, client, true, nil) + + if errDel := DeleteCachedSignatureRequired(context.Background(), testModelName, text); errDel != nil { + t.Fatalf("DeleteCachedSignatureRequired() error = %v", errDel) + } + if _, ok := client.values[signatureKVKey(testModelName, text)]; ok { + t.Fatalf("signature key was not deleted") + } + if client.delCount != 1 { + t.Fatalf("KVDel count = %d, want 1", client.delCount) + } +} + +func TestClearSignatureCacheHomeDoesNotPrefixDelete(t *testing.T) { + client := newFakeSignatureKVClient() + useFakeSignatureKVClient(t, client, true, nil) + + ClearSignatureCache("") + ClearSignatureCache(testModelName) + + if client.delCount != 0 { + t.Fatalf("ClearSignatureCache() KVDel count = %d, want 0", client.delCount) + } +} + +func TestGetCachedSignatureRequiredGeminiEmptyThinkingSentinel(t *testing.T) { + client := newFakeSignatureKVClient() + client.getErr = errors.New("get should not be called") + useFakeSignatureKVClient(t, client, true, nil) + + got, errGet := GetCachedSignatureRequired(context.Background(), "gemini-3-pro-preview", "") + if errGet != nil { + t.Fatalf("GetCachedSignatureRequired() error = %v", errGet) + } + if got != "skip_thought_signature_validator" { + t.Fatalf("GetCachedSignatureRequired() = %q, want Gemini sentinel", got) + } + if client.getCount != 0 { + t.Fatalf("KVGet count = %d, want 0", client.getCount) + } +} + func TestCacheSignature_DifferentModelGroups(t *testing.T) { ClearSignatureCache("") diff --git a/internal/home/client.go b/internal/home/client.go index fd7f98a2..a7ff8a5a 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -59,6 +59,13 @@ type clusterNodesEnvelope struct { Nodes []clusterNode `json:"nodes"` } +type KVSetOptions struct { + EX time.Duration + PX time.Duration + NX bool + XX bool +} + type Client struct { mu sync.Mutex @@ -531,6 +538,187 @@ func (c *Client) GetModels(ctx context.Context) ([]byte, error) { return raw, nil } +func buildKVSetArgs(key string, value []byte, opts KVSetOptions) ([]any, error) { + key = strings.TrimSpace(key) + if key == "" { + return nil, fmt.Errorf("home kv: key is empty") + } + if opts.EX > 0 && opts.PX > 0 { + return nil, fmt.Errorf("home kv: EX and PX are mutually exclusive") + } + if opts.EX < 0 || opts.PX < 0 { + return nil, fmt.Errorf("home kv: ttl must not be negative") + } + if opts.NX && opts.XX { + return nil, fmt.Errorf("home kv: NX and XX are mutually exclusive") + } + + args := []any{key, append([]byte(nil), value...)} + if opts.EX > 0 { + args = append(args, "EX", durationCeil(opts.EX, time.Second)) + } + if opts.PX > 0 { + args = append(args, "PX", durationCeil(opts.PX, time.Millisecond)) + } + if opts.NX { + args = append(args, "NX") + } + if opts.XX { + args = append(args, "XX") + } + return args, nil +} + +func durationCeil(value time.Duration, unit time.Duration) int64 { + if value <= 0 || unit <= 0 { + return 0 + } + return int64((value + unit - 1) / unit) +} + +func (c *Client) KVGet(ctx context.Context, key string) ([]byte, bool, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return nil, false, errClient + } + raw, errGet := cmd.Get(ctx, key).Bytes() + if errors.Is(errGet, redis.Nil) { + return nil, false, nil + } + if errGet != nil { + return nil, false, errGet + } + return append([]byte(nil), raw...), true, nil +} + +func (c *Client) KVSet(ctx context.Context, key string, value []byte, opts KVSetOptions) (bool, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return false, errClient + } + args, errArgs := buildKVSetArgs(key, value, opts) + if errArgs != nil { + return false, errArgs + } + result, errSet := cmd.Do(ctx, append([]any{"SET"}, args...)...).Result() + if errors.Is(errSet, redis.Nil) { + return false, nil + } + if errSet != nil { + return false, errSet + } + if result == nil { + return false, nil + } + return true, nil +} + +func (c *Client) KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) { + opts := KVSetOptions{NX: true} + if ttl > 0 { + opts.EX = ttl + } + return c.KVSet(ctx, key, value, opts) +} + +func (c *Client) KVDel(ctx context.Context, keys ...string) (int64, error) { + if len(keys) == 0 { + return 0, nil + } + cmd, errClient := c.commandClient() + if errClient != nil { + return 0, errClient + } + return cmd.Del(ctx, keys...).Result() +} + +func (c *Client) KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return false, errClient + } + return cmd.Expire(ctx, key, ttl).Result() +} + +func (c *Client) KVTTL(ctx context.Context, key string) (time.Duration, bool, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return 0, false, errClient + } + ttl, errTTL := cmd.TTL(ctx, key).Result() + if errTTL != nil { + return 0, false, errTTL + } + switch { + case ttl <= -2*time.Second: + return 0, false, nil + case ttl == -1*time.Second: + return 0, true, nil + default: + return ttl, true, nil + } +} + +func (c *Client) KVIncrBy(ctx context.Context, key string, delta int64) (int64, error) { + cmd, errClient := c.commandClient() + if errClient != nil { + return 0, errClient + } + return cmd.IncrBy(ctx, key, delta).Result() +} + +func (c *Client) KVMGet(ctx context.Context, keys ...string) ([][]byte, []bool, error) { + if len(keys) == 0 { + return nil, nil, nil + } + cmd, errClient := c.commandClient() + if errClient != nil { + return nil, nil, errClient + } + items, errMGet := cmd.MGet(ctx, keys...).Result() + if errMGet != nil { + return nil, nil, errMGet + } + values := make([][]byte, len(items)) + found := make([]bool, len(items)) + for i, item := range items { + switch typed := item.(type) { + case nil: + continue + case string: + values[i] = []byte(typed) + found[i] = true + case []byte: + values[i] = append([]byte(nil), typed...) + found[i] = true + default: + return nil, nil, fmt.Errorf("home kv: unsupported MGET item type %T", item) + } + } + return values, found, nil +} + +func (c *Client) KVMSet(ctx context.Context, pairs map[string][]byte) error { + if len(pairs) == 0 { + return nil + } + cmd, errClient := c.commandClient() + if errClient != nil { + return errClient + } + keys := make([]string, 0, len(pairs)) + for key := range pairs { + keys = append(keys, key) + } + sort.Strings(keys) + args := make([]any, 0, 1+len(keys)*2) + args = append(args, "MSET") + for _, key := range keys { + args = append(args, key, append([]byte(nil), pairs[key]...)) + } + return cmd.Do(ctx, args...).Err() +} + func headersToLowerMap(headers http.Header) map[string]string { if len(headers) == 0 { return nil diff --git a/internal/home/client_test.go b/internal/home/client_test.go index b0415d89..2a9f6789 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -1,12 +1,22 @@ package home import ( + "bufio" "context" "crypto/tls" "encoding/json" + "fmt" + "io" + "net" "net/http" + "reflect" + "strconv" + "strings" + "sync" "testing" + "time" + "github.com/redis/go-redis/v9" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" ) @@ -156,3 +166,236 @@ func TestFailoverAfterReconnectFailureDisabledDoesNotSwitchToClusterNode(t *test t.Fatalf("addr() = %q, want seed.example.com:8327", got) } } + +func TestBuildKVSetArgs(t *testing.T) { + args, errArgs := buildKVSetArgs("key", []byte("value"), KVSetOptions{EX: 2 * time.Second, NX: true}) + if errArgs != nil { + t.Fatalf("buildKVSetArgs(EX NX) error = %v", errArgs) + } + want := []any{"key", []byte("value"), "EX", int64(2), "NX"} + if !reflect.DeepEqual(args, want) { + t.Fatalf("buildKVSetArgs(EX NX) = %#v, want %#v", args, want) + } + + args, errArgs = buildKVSetArgs("key", []byte("value"), KVSetOptions{PX: 1500 * time.Millisecond, XX: true}) + if errArgs != nil { + t.Fatalf("buildKVSetArgs(PX XX) error = %v", errArgs) + } + want = []any{"key", []byte("value"), "PX", int64(1500), "XX"} + if !reflect.DeepEqual(args, want) { + t.Fatalf("buildKVSetArgs(PX XX) = %#v, want %#v", args, want) + } + + if _, errConflict := buildKVSetArgs("key", []byte("value"), KVSetOptions{EX: time.Second, PX: time.Millisecond}); errConflict == nil { + t.Fatalf("buildKVSetArgs(EX PX) error = nil, want error") + } + if _, errConflict := buildKVSetArgs("key", []byte("value"), KVSetOptions{NX: true, XX: true}); errConflict == nil { + t.Fatalf("buildKVSetArgs(NX XX) error = nil, want error") + } +} + +func TestKVGetConvertsRedisNilToMiss(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "GET") { + return "$-1\r\n" + } + return "-ERR unexpected command\r\n" + }) + + value, found, errGet := client.KVGet(context.Background(), "missing") + if errGet != nil { + t.Fatalf("KVGet() error = %v", errGet) + } + if found || value != nil { + t.Fatalf("KVGet() = %v, %v, want nil, false", value, found) + } +} + +func TestKVMGetConvertsNilItemsToMiss(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "MGET") { + return "*2\r\n$5\r\nvalue\r\n$-1\r\n" + } + return "-ERR unexpected command\r\n" + }) + + values, found, errMGet := client.KVMGet(context.Background(), "hit", "miss") + if errMGet != nil { + t.Fatalf("KVMGet() error = %v", errMGet) + } + if len(values) != 2 || len(found) != 2 { + t.Fatalf("KVMGet() lengths = %d, %d, want 2, 2", len(values), len(found)) + } + if !found[0] || string(values[0]) != "value" { + t.Fatalf("KVMGet()[0] = %q, %v, want value, true", values[0], found[0]) + } + if found[1] || values[1] != nil { + t.Fatalf("KVMGet()[1] = %v, %v, want nil, false", values[1], found[1]) + } +} + +func TestKVSetConditionUnmetReturnsFalse(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "SET") { + return "$-1\r\n" + } + return "-ERR unexpected command\r\n" + }) + + written, errSet := client.KVSet(context.Background(), "key", []byte("value"), KVSetOptions{NX: true}) + if errSet != nil { + t.Fatalf("KVSet() error = %v", errSet) + } + if written { + t.Fatalf("KVSet() written = true, want false") + } +} + +func TestKVMSetUsesStableKeyOrder(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "MSET") { + return "+OK\r\n" + } + return "-ERR unexpected command\r\n" + }) + + if errMSet := client.KVMSet(context.Background(), map[string][]byte{ + "b": []byte("2"), + "a": []byte("1"), + }); errMSet != nil { + t.Fatalf("KVMSet() error = %v", errMSet) + } + got := commands.Last() + want := []string{"MSET", "a", "1", "b", "2"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("MSET command = %#v, want %#v", got, want) + } +} + +type redisCommandLog struct { + mu sync.Mutex + commands [][]string +} + +func (l *redisCommandLog) Append(args []string) { + l.mu.Lock() + defer l.mu.Unlock() + l.commands = append(l.commands, append([]string(nil), args...)) +} + +func (l *redisCommandLog) Last() []string { + l.mu.Lock() + defer l.mu.Unlock() + if len(l.commands) == 0 { + return nil + } + return append([]string(nil), l.commands[len(l.commands)-1]...) +} + +func newRedisCommandTestClient(t *testing.T, handler func([]string) string) (*Client, *redisCommandLog) { + t.Helper() + + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + log := &redisCommandLog{} + done := make(chan struct{}) + go func() { + defer close(done) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRedisCommandTestConn(conn, log, handler) + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-done + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener addr: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse listener port: %v", errPort) + } + client := New(config.HomeConfig{ + Enabled: true, + Host: host, + Port: port, + DisableClusterDiscovery: true, + }) + client.cmd = redis.NewClient(&redis.Options{ + Addr: listener.Addr().String(), + Protocol: 2, + DisableIdentity: true, + MaxRetries: -1, + ContextTimeoutEnabled: true, + }) + t.Cleanup(func() { + client.Close() + }) + return client, log +} + +func serveRedisCommandTestConn(conn net.Conn, log *redisCommandLog, handler func([]string) string) { + defer func() { + _ = conn.Close() + }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRedisCommand(reader) + if errRead != nil { + return + } + log.Append(args) + response := "+OK\r\n" + if handler != nil { + response = handler(args) + } + if _, errWrite := io.WriteString(conn, response); errWrite != nil { + return + } + } +} + +func readRedisCommand(reader *bufio.Reader) ([]string, error) { + line, errRead := reader.ReadString('\n') + if errRead != nil { + return nil, errRead + } + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "*") { + return nil, fmt.Errorf("expected array, got %q", line) + } + count, errCount := strconv.Atoi(strings.TrimPrefix(line, "*")) + if errCount != nil { + return nil, errCount + } + args := make([]string, 0, count) + for i := 0; i < count; i++ { + bulkLine, errBulk := reader.ReadString('\n') + if errBulk != nil { + return nil, errBulk + } + bulkLine = strings.TrimSpace(bulkLine) + if !strings.HasPrefix(bulkLine, "$") { + return nil, fmt.Errorf("expected bulk string, got %q", bulkLine) + } + size, errSize := strconv.Atoi(strings.TrimPrefix(bulkLine, "$")) + if errSize != nil { + return nil, errSize + } + payload := make([]byte, size+2) + if _, errFull := io.ReadFull(reader, payload); errFull != nil { + return nil, errFull + } + args = append(args, string(payload[:size])) + } + return args, nil +} diff --git a/internal/home/kv_helpers.go b/internal/home/kv_helpers.go new file mode 100644 index 00000000..7ca21700 --- /dev/null +++ b/internal/home/kv_helpers.go @@ -0,0 +1,189 @@ +package home + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "strings" + "time" + + log "github.com/sirupsen/logrus" +) + +func HashKeyPart(value string) string { + sum := sha256.Sum256([]byte(value)) + return hex.EncodeToString(sum[:]) +} + +func CurrentKVClient() (*Client, bool, error) { + client := Current() + if client == nil { + return nil, false, nil + } + if !client.Enabled() { + return nil, true, fmt.Errorf("home kv store unavailable: %w", ErrDisabled) + } + if !client.HeartbeatOK() { + return nil, true, fmt.Errorf("home kv store unavailable: %w", ErrNotConnected) + } + return client, true, nil +} + +func KVGetJSONRequired(ctx context.Context, key string, out any) (bool, bool, error) { + client, homeMode, errClient := CurrentKVClient() + if !homeMode || errClient != nil { + return homeMode, false, errClient + } + raw, found, errGet := client.KVGet(ctx, key) + if errGet != nil || !found { + return true, false, errGet + } + if errUnmarshal := json.Unmarshal(raw, out); errUnmarshal != nil { + return true, false, errUnmarshal + } + return true, true, nil +} + +func KVSetJSONRequired(ctx context.Context, key string, value any, ttl time.Duration) (bool, error) { + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + return false, errMarshal + } + return KVSetBytesRequired(ctx, key, raw, ttl) +} + +func KVSetBytesRequired(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) { + client, homeMode, errClient := CurrentKVClient() + if !homeMode || errClient != nil { + return homeMode, errClient + } + written, errSet := client.KVSet(ctx, key, value, kvSetOptionsForTTL(ttl)) + if errSet != nil { + return true, errSet + } + if !written { + return true, fmt.Errorf("home kv store unavailable") + } + return true, nil +} + +func KVSetNXRequired(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, bool, error) { + client, homeMode, errClient := CurrentKVClient() + if !homeMode || errClient != nil { + return homeMode, false, errClient + } + written, errSet := client.KVSetNX(ctx, key, value, ttl) + return true, written, errSet +} + +func KVDelRequired(ctx context.Context, keys ...string) (bool, int64, error) { + client, homeMode, errClient := CurrentKVClient() + if !homeMode || errClient != nil { + return homeMode, 0, errClient + } + deleted, errDel := client.KVDel(ctx, keys...) + return true, deleted, errDel +} + +func KVExpireRequired(ctx context.Context, key string, ttl time.Duration) (bool, error) { + client, homeMode, errClient := CurrentKVClient() + if !homeMode || errClient != nil { + return homeMode, errClient + } + _, errExpire := client.KVExpire(ctx, key, ttl) + return true, errExpire +} + +func KVGetJSONBestEffort(ctx context.Context, key string, out any) (bool, bool) { + homeMode, found, errGet := KVGetJSONRequired(ctx, key, out) + if errGet != nil { + log.Errorf("home kv best-effort get failed prefix=%s: %v", kvLogPrefix(key), errGet) + return homeMode, false + } + return homeMode, found +} + +func KVSetJSONBestEffort(ctx context.Context, key string, value any, ttl time.Duration) bool { + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + log.Errorf("home kv best-effort set failed prefix=%s: %v", kvLogPrefix(key), errMarshal) + return false + } + return KVSetBytesBestEffort(ctx, key, raw, ttl) +} + +func KVSetBytesBestEffort(ctx context.Context, key string, value []byte, ttl time.Duration) bool { + homeMode, errSet := KVSetBytesRequired(ctx, key, value, ttl) + if !homeMode { + return false + } + if errSet != nil { + log.Errorf("home kv best-effort set failed prefix=%s: %v", kvLogPrefix(key), errSet) + return false + } + return true +} + +func KVSetNXBestEffort(ctx context.Context, key string, value []byte, ttl time.Duration) bool { + homeMode, written, errSet := KVSetNXRequired(ctx, key, value, ttl) + if !homeMode { + return false + } + if errSet != nil { + log.Errorf("home kv best-effort setnx failed prefix=%s: %v", kvLogPrefix(key), errSet) + return false + } + return written +} + +func KVDelBestEffort(ctx context.Context, keys ...string) bool { + homeMode, _, errDel := KVDelRequired(ctx, keys...) + if !homeMode { + return false + } + if errDel != nil { + log.Errorf("home kv best-effort del failed prefix=%s: %v", kvLogPrefix(firstKVKey(keys)), errDel) + return false + } + return true +} + +func KVExpireBestEffort(ctx context.Context, key string, ttl time.Duration) bool { + homeMode, errExpire := KVExpireRequired(ctx, key, ttl) + if !homeMode { + return false + } + if errExpire != nil { + log.Errorf("home kv best-effort expire failed prefix=%s: %v", kvLogPrefix(key), errExpire) + return false + } + return true +} + +func kvSetOptionsForTTL(ttl time.Duration) KVSetOptions { + if ttl <= 0 { + return KVSetOptions{} + } + return KVSetOptions{EX: ttl} +} + +func kvLogPrefix(key string) string { + key = strings.TrimSpace(key) + if key == "" { + return "unknown" + } + parts := strings.Split(key, ":") + if len(parts) >= 2 { + return parts[0] + ":" + parts[1] + ":*" + } + return parts[0] + ":*" +} + +func firstKVKey(keys []string) string { + if len(keys) == 0 { + return "" + } + return keys[0] +} diff --git a/internal/home/kv_helpers_test.go b/internal/home/kv_helpers_test.go new file mode 100644 index 00000000..012d377a --- /dev/null +++ b/internal/home/kv_helpers_test.go @@ -0,0 +1,110 @@ +package home + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + log "github.com/sirupsen/logrus" +) + +func TestHashKeyPart(t *testing.T) { + first := HashKeyPart("secret-value") + again := HashKeyPart("secret-value") + other := HashKeyPart("other-value") + if first == "" || len(first) != 64 { + t.Fatalf("HashKeyPart() = %q, want 64 hex chars", first) + } + if first != again { + t.Fatalf("HashKeyPart() is not stable") + } + if first == other { + t.Fatalf("HashKeyPart() returned same hash for different inputs") + } + if strings.Contains(first, "secret") || strings.Contains(first, "value") { + t.Fatalf("HashKeyPart() leaked input: %q", first) + } +} + +func TestKVRequiredHelpersReturnNonHomeMode(t *testing.T) { + ClearCurrent() + t.Cleanup(ClearCurrent) + + var out map[string]string + homeMode, found, errGet := KVGetJSONRequired(context.Background(), "key", &out) + if errGet != nil { + t.Fatalf("KVGetJSONRequired() error = %v", errGet) + } + if homeMode || found { + t.Fatalf("KVGetJSONRequired() = homeMode %v found %v, want false false", homeMode, found) + } +} + +func TestCurrentKVClientUnavailableErrors(t *testing.T) { + t.Cleanup(ClearCurrent) + + disabled := New(config.HomeConfig{Enabled: false}) + SetCurrent(disabled) + if _, homeMode, errClient := CurrentKVClient(); !homeMode || errClient == nil { + t.Fatalf("CurrentKVClient(disabled) = homeMode %v err %v, want true error", homeMode, errClient) + } + + notReady := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 1}) + SetCurrent(notReady) + if _, homeMode, errClient := CurrentKVClient(); !homeMode || errClient == nil { + t.Fatalf("CurrentKVClient(no heartbeat) = homeMode %v err %v, want true error", homeMode, errClient) + } +} + +func TestKVRequiredHelpersPropagateClientErrors(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + return "-ERR home kv unavailable\r\n" + }) + client.heartbeatOK.Store(true) + SetCurrent(client) + t.Cleanup(ClearCurrent) + + var out map[string]string + homeMode, _, errGet := KVGetJSONRequired(context.Background(), "cpa:test:key", &out) + if !homeMode || errGet == nil { + t.Fatalf("KVGetJSONRequired() = homeMode %v err %v, want true error", homeMode, errGet) + } + homeMode, errSet := KVSetJSONRequired(context.Background(), "cpa:test:key", map[string]string{"value": "secret"}, 0) + if !homeMode || errSet == nil { + t.Fatalf("KVSetJSONRequired() = homeMode %v err %v, want true error", homeMode, errSet) + } +} + +func TestKVBestEffortWriteSwallowsErrorAndRedactsLog(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + return "-ERR home kv unavailable\r\n" + }) + client.heartbeatOK.Store(true) + SetCurrent(client) + t.Cleanup(ClearCurrent) + + logger := log.StandardLogger() + previousOutput := logger.Out + previousLevel := log.GetLevel() + buffer := &bytes.Buffer{} + log.SetOutput(buffer) + log.SetLevel(log.ErrorLevel) + t.Cleanup(func() { + log.SetOutput(previousOutput) + log.SetLevel(previousLevel) + }) + + ok := KVSetJSONBestEffort(context.Background(), "cpa:test:secret-key", map[string]string{"value": "secret-value"}, 0) + if ok { + t.Fatalf("KVSetJSONBestEffort() = true, want false") + } + logText := buffer.String() + if !strings.Contains(logText, "cpa:test:*") { + t.Fatalf("log = %q, want redacted key prefix", logText) + } + if strings.Contains(logText, "secret-key") || strings.Contains(logText, "secret-value") { + t.Fatalf("log leaked key or value: %q", logText) + } +} diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index cd3b191c..3ce78079 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -25,6 +25,7 @@ import ( "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" @@ -100,6 +101,17 @@ var ( } ) +type antigravityKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) + KVDel(ctx context.Context, keys ...string) (int64, error) +} + +var currentAntigravityKVClient = func() (antigravityKVClient, bool, error) { + return homekv.CurrentKVClient() +} + type antigravityCreditsBalance struct { CreditAmount float64 MinCreditAmount float64 @@ -120,26 +132,62 @@ type antigravityTokenRefreshData struct { } func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool { - if auth == nil || strings.TrimSpace(auth.ID) == "" { + ok, err := antigravityAuthHasCreditsRequired(context.Background(), auth) + if err != nil { + log.Errorf("antigravity executor: home kv credits check error: %v", err) return false } - if hint, ok := cliproxyauth.GetAntigravityCreditsHint(auth.ID); ok && hint.Known { - return hint.Available + return ok +} + +func antigravityAuthHasCreditsRequired(ctx context.Context, auth *cliproxyauth.Auth) (bool, error) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return false, nil } - val, ok := antigravityCreditsBalanceByAuth.Load(strings.TrimSpace(auth.ID)) + authID := strings.TrimSpace(auth.ID) + if hint, ok, errHint := cliproxyauth.GetAntigravityCreditsHintRequired(ctx, authID); errHint != nil { + return false, errHint + } else if ok && hint.Known { + return hint.Available, nil + } + + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return false, errClient + } + raw, found, errBalance := client.KVGet(ctx, antigravityCreditsBalanceKey(authID)) + if errBalance != nil { + return false, errBalance + } + if !found { + return true, nil + } + var homeBalance antigravityCreditsBalance + if errUnmarshal := json.Unmarshal(raw, &homeBalance); errUnmarshal != nil { + return false, errUnmarshal + } + return antigravityCreditsBalanceAvailable(authID, homeBalance), nil + } + + val, ok := antigravityCreditsBalanceByAuth.Load(authID) if !ok { - return true // optimistic: assume credits available when balance unknown + return true, nil // optimistic: assume credits available when balance unknown } bal, valid := val.(antigravityCreditsBalance) if !valid { - antigravityCreditsBalanceByAuth.Delete(strings.TrimSpace(auth.ID)) - return false + antigravityCreditsBalanceByAuth.Delete(authID) + return false, nil } + return antigravityCreditsBalanceAvailable(authID, bal), nil +} + +func antigravityCreditsBalanceAvailable(authID string, bal antigravityCreditsBalance) bool { if !bal.Known { return false } available := bal.CreditAmount >= bal.MinCreditAmount - cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(auth.ID), cliproxyauth.AntigravityCreditsHint{ + cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(authID), cliproxyauth.AntigravityCreditsHint{ Known: true, Available: available, CreditAmount: bal.CreditAmount, @@ -249,7 +297,7 @@ func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cli return client } -func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []byte) ([]byte, error) { +func validateAntigravityRequestSignatures(ctx context.Context, modelName string, from sdktranslator.Format, rawJSON []byte) ([]byte, error) { if from.String() != "claude" { return rawJSON, nil } @@ -258,6 +306,9 @@ func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []b rawJSON = antigravityclaude.StripEmptySignatureThinkingBlocks(rawJSON) logAntigravitySignatureStrip(before, countClaudeThinkingBlocks(rawJSON), "prefix_cleanup", "empty_or_non_claude_signature") if cache.SignatureCacheEnabled() { + if errRequire := antigravityclaude.RequireCachedThinkingSignatures(ctx, modelName, rawJSON); errRequire != nil { + return nil, homeKVUnavailableStatusErr(errRequire) + } return rawJSON, nil } if !cache.SignatureBypassStrictMode() { @@ -511,11 +562,12 @@ func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) { ExplicitBalanceExhausted: true, } antigravityCreditsFailureByAuth.Store(authID, state) - antigravityCreditsBalanceByAuth.Store(authID, antigravityCreditsBalance{ + bal := antigravityCreditsBalance{ CreditAmount: 0, MinCreditAmount: 1, Known: true, - }) + } + storeAntigravityCreditsBalanceBestEffort(authID, bal) cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ Known: true, Available: false, @@ -568,7 +620,9 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} } baseModel := thinking.ParseSuffix(req.Model).ModelName - if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return resp, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) d := remaining return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} @@ -591,7 +645,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) if errValidate != nil { return resp, errValidate } @@ -689,7 +743,10 @@ attemptLoop: } case antigravity429DecisionShortCooldownSwitchAuth: if decision.retryAfter != nil && *decision.retryAfter > 0 { - markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) } case antigravity429DecisionFullQuotaExhausted: @@ -775,7 +832,9 @@ attemptLoop: // executeClaudeNonStream performs a claude non-streaming request to the Antigravity API. func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { baseModel := thinking.ParseSuffix(req.Model).ModelName - if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return resp, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) d := remaining return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} @@ -793,7 +852,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) if errValidate != nil { return resp, errValidate } @@ -906,7 +965,10 @@ attemptLoop: } case antigravity429DecisionShortCooldownSwitchAuth: if decision.retryAfter != nil && *decision.retryAfter > 0 { - markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) } case antigravity429DecisionFullQuotaExhausted: @@ -1239,7 +1301,9 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya baseModel := thinking.ParseSuffix(req.Model).ModelName ctx = context.WithValue(ctx, "alt", "") - if inCooldown, remaining := antigravityIsInShortCooldown(auth, baseModel, time.Now()); inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return nil, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) d := remaining return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} @@ -1257,7 +1321,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya originalPayloadSource = opts.OriginalRequest } originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(from, originalPayload) + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) if errValidate != nil { return nil, errValidate } @@ -1370,7 +1434,10 @@ attemptLoop: } case antigravity429DecisionShortCooldownSwitchAuth: if decision.retryAfter != nil && *decision.retryAfter > 0 { - markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter) + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return nil, err + } log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) } case antigravity429DecisionFullQuotaExhausted: @@ -1564,7 +1631,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut if len(opts.OriginalRequest) > 0 { originalPayloadSource = opts.OriginalRequest } - originalPayloadSource, errValidate := validateAntigravityRequestSignatures(from, originalPayloadSource) + originalPayloadSource, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayloadSource) if errValidate != nil { return cliproxyexecutor.Response{}, errValidate } @@ -1770,6 +1837,34 @@ func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Con return } + if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { + if errClient != nil { + log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errClient) + return + } + written, errSetNX := client.KVSetNX(context.Background(), antigravityCreditsRefreshLockKey(authID), []byte("1"), antigravityCreditsHintRefreshInterval) + if errSetNX != nil { + log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errSetNX) + return + } + if !written { + return + } + refreshCtx := context.Background() + if ctx != nil { + if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { + refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) + } + } + refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) + authCopy := auth.Clone() + go func(auth *cliproxyauth.Auth, token string) { + defer cancel() + e.updateAntigravityCreditsBalance(refreshCtx, auth, token) + }(authCopy, accessToken) + return + } + state := &antigravityCreditsHintRefreshState{} if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded { if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil { @@ -2048,7 +2143,7 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Contex PaidTierID: paidTierID, Known: true, } - antigravityCreditsBalanceByAuth.Store(authID, bal) + storeAntigravityCreditsBalanceBestEffort(authID, bal) cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ Known: true, Available: creditAmount >= minAmount, @@ -2406,34 +2501,147 @@ func antigravityShortCooldownKey(auth *cliproxyauth.Auth, modelName string) stri return authID + "|" + modelName + "|sc" } +func antigravityCreditsBalanceKey(authID string) string { + return "cpa:antigravity:credits-balance:" + strings.TrimSpace(authID) +} + +func antigravityCreditsRefreshLockKey(authID string) string { + return "cpa:antigravity:credits-refresh-lock:" + strings.TrimSpace(authID) +} + +func antigravityShortCooldownKVKey(auth *cliproxyauth.Auth, modelName string) string { + if auth == nil { + return "" + } + authID := strings.TrimSpace(auth.ID) + modelName = strings.TrimSpace(modelName) + if authID == "" || modelName == "" { + return "" + } + return "cpa:antigravity:short-cooldown:" + authID + ":" + homekv.HashKeyPart(modelName) +} + func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) { + inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now) + if errCooldown != nil { + log.Errorf("antigravity executor: home kv cooldown read error: %v", errCooldown) + return false, 0 + } + return inCooldown, remaining +} + +func antigravityIsInShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration, error) { + kvKey := antigravityShortCooldownKVKey(auth, modelName) + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return false, 0, errClient + } + if kvKey == "" { + return false, 0, nil + } + raw, found, errGet := client.KVGet(ctx, kvKey) + if errGet != nil || !found { + return false, 0, errGet + } + untilNano, errParse := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64) + if errParse != nil { + return false, 0, errParse + } + remaining := time.Unix(0, untilNano).Sub(now) + if remaining <= 0 { + if _, errDel := client.KVDel(ctx, kvKey); errDel != nil { + return false, 0, errDel + } + return false, 0, nil + } + return true, remaining, nil + } + key := antigravityShortCooldownKey(auth, modelName) if key == "" { - return false, 0 + return false, 0, nil } value, ok := antigravityShortCooldownByAuth.Load(key) if !ok { - return false, 0 + return false, 0, nil } until, ok := value.(time.Time) if !ok || until.IsZero() { antigravityShortCooldownByAuth.Delete(key) - return false, 0 + return false, 0, nil } remaining := until.Sub(now) if remaining <= 0 { antigravityShortCooldownByAuth.Delete(key) - return false, 0 + return false, 0, nil } - return true, remaining + return true, remaining, nil } func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) { + if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil { + log.Errorf("antigravity executor: home kv cooldown write error: %v", errMark) + } +} + +func markAntigravityShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) error { + kvKey := antigravityShortCooldownKVKey(auth, modelName) + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return errClient + } + if kvKey == "" || duration <= 0 { + return nil + } + until := now.Add(duration) + written, errSet := client.KVSet(ctx, kvKey, []byte(strconv.FormatInt(until.UnixNano(), 10)), homekv.KVSetOptions{EX: duration + 5*time.Second}) + if errSet != nil { + return errSet + } + if !written { + return fmt.Errorf("home kv store unavailable") + } + return nil + } + key := antigravityShortCooldownKey(auth, modelName) if key == "" { - return + return nil } antigravityShortCooldownByAuth.Store(key, now.Add(duration)) + return nil +} + +func storeAntigravityCreditsBalanceBestEffort(authID string, bal antigravityCreditsBalance) { + authID = strings.TrimSpace(authID) + if authID == "" { + return + } + if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { + if errClient != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errClient) + return + } + raw, errMarshal := json.Marshal(bal) + if errMarshal != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errMarshal) + return + } + if _, errSet := client.KVSet(context.Background(), antigravityCreditsBalanceKey(authID), raw, homekv.KVSetOptions{EX: 30 * time.Minute}); errSet != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errSet) + } + return + } + antigravityCreditsBalanceByAuth.Store(authID, bal) +} + +func homeKVUnavailableStatusErr(cause error) statusErr { + if cause == nil { + return statusErr{code: http.StatusServiceUnavailable, msg: "home kv store unavailable"} + } + return statusErr{code: http.StatusServiceUnavailable, msg: fmt.Sprintf("home kv store unavailable: %v", cause)} } func antigravityNoCapacityRetryDelay(attempt int) time.Duration { diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index ac523339..507a57b3 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -2,6 +2,8 @@ package executor import ( "context" + "encoding/json" + "errors" "io" "net/http" "net/http/httptest" @@ -11,6 +13,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" @@ -23,6 +26,105 @@ func resetAntigravityCreditsRetryState() { antigravityCreditsHintRefreshByID = sync.Map{} } +type fakeAntigravityKVClient struct { + values map[string][]byte + getErr error + setErr error + setNXErr error + delErr error + setNXResult bool + getCount int + setCount int + setNXCount int + delCount int + lastSetTTL time.Duration + lastSetNXTTL time.Duration + lastSetNXKey string + lastSetKey string +} + +func newFakeAntigravityKVClient() *fakeAntigravityKVClient { + return &fakeAntigravityKVClient{ + values: make(map[string][]byte), + setNXResult: true, + } +} + +func (c *fakeAntigravityKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeAntigravityKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.setCount++ + c.lastSetKey = key + c.lastSetTTL = opts.EX + if c.setErr != nil { + return false, c.setErr + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + +func (c *fakeAntigravityKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) { + c.setNXCount++ + c.lastSetNXKey = key + c.lastSetNXTTL = ttl + if c.setNXErr != nil { + return false, c.setNXErr + } + if _, ok := c.values[key]; ok { + return false, nil + } + if c.setNXResult { + c.values[key] = append([]byte(nil), value...) + return true, nil + } + return false, nil +} + +func (c *fakeAntigravityKVClient) KVDel(_ context.Context, keys ...string) (int64, error) { + c.delCount++ + if c.delErr != nil { + return 0, c.delErr + } + var deleted int64 + for _, key := range keys { + if _, ok := c.values[key]; ok { + delete(c.values, key) + deleted++ + } + } + return deleted, nil +} + +func useFakeAntigravityKVClient(t *testing.T, client *fakeAntigravityKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentAntigravityKVClient + currentAntigravityKVClient = func() (antigravityKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentAntigravityKVClient = previous + }) +} + +func mustAntigravityJSON(t *testing.T, value any) []byte { + t.Helper() + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + t.Fatalf("marshal value: %v", errMarshal) + } + return raw +} + func TestClassifyAntigravity429(t *testing.T) { t.Run("quota exhausted", func(t *testing.T) { body := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`) @@ -379,6 +481,134 @@ func TestAntigravityAuthHasCredits(t *testing.T) { }) } +func TestAntigravityAuthHasCreditsRequiredHomeBalanceUsesKV(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + const authID = "home-balance-auth" + client := newFakeAntigravityKVClient() + client.values[antigravityCreditsBalanceKey(authID)] = mustAntigravityJSON(t, antigravityCreditsBalance{ + CreditAmount: 10, + MinCreditAmount: 50, + Known: true, + }) + useFakeAntigravityKVClient(t, client, true, nil) + antigravityCreditsBalanceByAuth.Store(authID, antigravityCreditsBalance{ + CreditAmount: 25000, + MinCreditAmount: 50, + Known: true, + }) + + ok, errCredits := antigravityAuthHasCreditsRequired(context.Background(), &cliproxyauth.Auth{ID: authID}) + if errCredits != nil { + t.Fatalf("antigravityAuthHasCreditsRequired() error = %v", errCredits) + } + if ok { + t.Fatalf("antigravityAuthHasCreditsRequired() = true, want Home KV balance to win over local cache") + } + if client.getCount != 1 { + t.Fatalf("KVGet count = %d, want 1", client.getCount) + } +} + +func TestStoreAntigravityCreditsBalanceBestEffortHomeKV(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + const authID = "home-balance-write-auth" + client := newFakeAntigravityKVClient() + useFakeAntigravityKVClient(t, client, true, nil) + + storeAntigravityCreditsBalanceBestEffort(authID, antigravityCreditsBalance{ + CreditAmount: 25000, + MinCreditAmount: 50, + Known: true, + }) + + if client.setCount != 1 || client.lastSetKey != antigravityCreditsBalanceKey(authID) || client.lastSetTTL != 30*time.Minute { + t.Fatalf("KVSet count/key/ttl = %d/%s/%v, want 1/%s/30m", client.setCount, client.lastSetKey, client.lastSetTTL, antigravityCreditsBalanceKey(authID)) + } + if _, ok := antigravityCreditsBalanceByAuth.Load(authID); ok { + t.Fatalf("local balance cache was populated in Home mode") + } +} + +func TestAntigravityShortCooldownRequiredHomeKV(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + client := newFakeAntigravityKVClient() + useFakeAntigravityKVClient(t, client, true, nil) + auth := &cliproxyauth.Auth{ID: "home-cooldown-auth"} + now := time.Now() + duration := 30 * time.Second + + if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, "claude-sonnet-4-5", now, duration); errMark != nil { + t.Fatalf("markAntigravityShortCooldownRequired() error = %v", errMark) + } + if client.setCount != 1 || client.lastSetTTL != duration+5*time.Second { + t.Fatalf("KVSet count/ttl = %d/%v, want 1/%v", client.setCount, client.lastSetTTL, duration+5*time.Second) + } + antigravityShortCooldownByAuth = sync.Map{} + inCooldown, remaining, errRead := antigravityIsInShortCooldownRequired(context.Background(), auth, "claude-sonnet-4-5", now.Add(5*time.Second)) + if errRead != nil { + t.Fatalf("antigravityIsInShortCooldownRequired() error = %v", errRead) + } + if !inCooldown || remaining <= 0 { + t.Fatalf("cooldown = %v remaining %v, want active Home KV cooldown", inCooldown, remaining) + } +} + +func TestAntigravityShortCooldownRequiredHomeKVFailures(t *testing.T) { + auth := &cliproxyauth.Auth{ID: "home-cooldown-failure-auth"} + for _, tc := range []struct { + name string + client *fakeAntigravityKVClient + write bool + }{ + {name: "read", client: &fakeAntigravityKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "write", client: &fakeAntigravityKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}, write: true}, + {name: "delete-expired", client: &fakeAntigravityKVClient{ + values: map[string][]byte{ + antigravityShortCooldownKVKey(auth, "claude-sonnet-4-5"): []byte("1"), + }, + delErr: errors.New("delete failed"), + }}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeAntigravityKVClient(t, tc.client, true, nil) + if tc.write { + if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, "claude-sonnet-4-5", time.Now(), time.Second); errMark == nil { + t.Fatalf("markAntigravityShortCooldownRequired() error = nil, want error") + } + return + } + if _, _, errRead := antigravityIsInShortCooldownRequired(context.Background(), auth, "claude-sonnet-4-5", time.Now()); errRead == nil { + t.Fatalf("antigravityIsInShortCooldownRequired() error = nil, want error") + } + }) + } +} + +func TestMaybeRefreshAntigravityCreditsHintHomeRefreshThrottleUsesSetNX(t *testing.T) { + resetAntigravityCreditsRetryState() + t.Cleanup(resetAntigravityCreditsRetryState) + client := newFakeAntigravityKVClient() + client.setNXResult = false + useFakeAntigravityKVClient(t, client, true, nil) + exec := NewAntigravityExecutor(&config.Config{ + QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, + }) + auth := &cliproxyauth.Auth{ID: "home-refresh-throttle-auth"} + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) { + t.Fatalf("refresh request should not run when Home KV throttle lock is not acquired") + return nil, nil + })) + + exec.maybeRefreshAntigravityCreditsHint(ctx, auth, "access-token") + + if client.setNXCount != 1 || client.lastSetNXKey != antigravityCreditsRefreshLockKey(auth.ID) || client.lastSetNXTTL != antigravityCreditsHintRefreshInterval { + t.Fatalf("KVSetNX count/key/ttl = %d/%s/%v, want 1/%s/%v", client.setNXCount, client.lastSetNXKey, client.lastSetNXTTL, antigravityCreditsRefreshLockKey(auth.ID), antigravityCreditsHintRefreshInterval) + } +} + type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { diff --git a/internal/runtime/executor/antigravity_executor_signature_test.go b/internal/runtime/executor/antigravity_executor_signature_test.go index 8383614d..c35190e4 100644 --- a/internal/runtime/executor/antigravity_executor_signature_test.go +++ b/internal/runtime/executor/antigravity_executor_signature_test.go @@ -101,7 +101,7 @@ func TestAntigravityExecutor_StrictBypassStripsInvalidSignature(t *testing.T) { payload := invalidClaudeThinkingPayload() from := sdktranslator.FromString("claude") - output, err := validateAntigravityRequestSignatures(from, payload) + output, err := validateAntigravityRequestSignatures(context.Background(), "claude-sonnet-4-5-thinking", from, payload) if err != nil { t.Fatalf("strict bypass should strip invalid signatures instead of rejecting request: %v", err) } @@ -140,7 +140,7 @@ func TestAntigravityExecutor_StrictBypassLogsStrippedInvalidSignature(t *testing }`) from := sdktranslator.FromString("claude") - if _, err := validateAntigravityRequestSignatures(from, payload); err != nil { + if _, err := validateAntigravityRequestSignatures(context.Background(), "claude-sonnet-4-5-thinking", from, payload); err != nil { t.Fatalf("strict bypass should strip invalid signatures instead of rejecting request: %v", err) } @@ -229,7 +229,7 @@ func TestAntigravityExecutor_NonStrictBypassSkipsPrecheck(t *testing.T) { payload := invalidClaudeThinkingPayload() from := sdktranslator.FromString("claude") - _, err := validateAntigravityRequestSignatures(from, payload) + _, err := validateAntigravityRequestSignatures(context.Background(), "claude-sonnet-4-5-thinking", from, payload) if err != nil { t.Fatalf("non-strict bypass should skip precheck, got: %v", err) } @@ -245,7 +245,7 @@ func TestAntigravityExecutor_CacheModeSkipsPrecheck(t *testing.T) { payload := invalidClaudeThinkingPayload() from := sdktranslator.FromString("claude") - _, err := validateAntigravityRequestSignatures(from, payload) + _, err := validateAntigravityRequestSignatures(context.Background(), "claude-sonnet-4-5-thinking", from, payload) if err != nil { t.Fatalf("cache mode should skip precheck, got: %v", err) } diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index 22de9183..dd5933a9 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -193,7 +193,10 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) // based on client type and configuration. - body = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + if err != nil { + return resp, err + } requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) @@ -241,7 +244,9 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r if err != nil { return resp, err } - applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg) + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg); errHeaders != nil { + return resp, errHeaders + } var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -375,7 +380,10 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) // based on client type and configuration. - body = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + if err != nil { + return nil, err + } requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) @@ -419,7 +427,9 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if err != nil { return nil, err } - applyClaudeHeaders(httpReq, auth, apiKey, true, extraBetas, e.cfg) + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, true, extraBetas, e.cfg); errHeaders != nil { + return nil, errHeaders + } var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -657,7 +667,9 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut if err != nil { return cliproxyexecutor.Response{}, err } - applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg) + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg); errHeaders != nil { + return cliproxyexecutor.Response{}, errHeaders + } var authID, authLabel, authType, authValue string if auth != nil { authID = auth.ID @@ -956,7 +968,10 @@ func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadClos return body, nil } -func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string, cfg *config.Config) { +func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string, cfg *config.Config) error { + if r == nil { + return nil + } hdrDefault := func(cfgVal, fallback string) string { if cfgVal != "" { return cfgVal @@ -986,7 +1001,11 @@ func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stabilizeDeviceProfile := helps.ClaudeDeviceProfileStabilizationEnabled(cfg) var deviceProfile helps.ClaudeDeviceProfile if stabilizeDeviceProfile { - deviceProfile = helps.ResolveClaudeDeviceProfile(auth, apiKey, ginHeaders, cfg) + var errDeviceProfile error + deviceProfile, errDeviceProfile = helps.ResolveClaudeDeviceProfileRequired(r.Context(), auth, apiKey, ginHeaders, cfg) + if errDeviceProfile != nil { + return errDeviceProfile + } } baseBetas := "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05,structured-outputs-2025-12-15,fast-mode-2026-02-01,redact-thinking-2026-02-12,token-efficient-tools-2026-03-28" @@ -1031,7 +1050,11 @@ func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Lang", "js") misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Timeout", hdrDefault(hd.Timeout, "600")) // Session ID: stable per auth/apiKey, matches Claude Code's X-Claude-Code-Session-Id header. - misc.EnsureHeader(r.Header, ginHeaders, "X-Claude-Code-Session-Id", helps.CachedSessionID(apiKey)) + sessionID, errSessionID := helps.CachedSessionIDRequired(r.Context(), apiKey) + if errSessionID != nil { + return errSessionID + } + misc.EnsureHeader(r.Header, ginHeaders, "X-Claude-Code-Session-Id", sessionID) // Per-request UUID, matches Claude Code's x-client-request-id for first-party API. if isAnthropicBase { misc.EnsureHeader(r.Header, ginHeaders, "x-client-request-id", uuid.New().String()) @@ -1066,6 +1089,7 @@ func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, if stream { r.Header.Set("Accept-Encoding", "identity") } + return nil } func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { @@ -1592,25 +1616,33 @@ func getCloakConfigFromAuth(auth *cliproxyauth.Auth) (string, bool, []string, bo // injectFakeUserID generates and injects a fake user ID into the request metadata. // When useCache is false, a new user ID is generated for every call. -func injectFakeUserID(payload []byte, apiKey string, useCache bool) []byte { - generateID := func() string { +func injectFakeUserID(ctx context.Context, payload []byte, apiKey string, useCache bool) ([]byte, error) { + generateID := func() (string, error) { if useCache { - return helps.CachedUserID(apiKey) + return helps.CachedUserIDRequired(ctx, apiKey) } - return helps.GenerateFakeUserID() + return helps.GenerateFakeUserID(), nil } metadata := gjson.GetBytes(payload, "metadata") if !metadata.Exists() { - payload, _ = sjson.SetBytes(payload, "metadata.user_id", generateID()) - return payload + userID, errUserID := generateID() + if errUserID != nil { + return nil, errUserID + } + payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) + return payload, nil } existingUserID := gjson.GetBytes(payload, "metadata.user_id").String() if existingUserID == "" || !helps.IsValidUserID(existingUserID) { - payload, _ = sjson.SetBytes(payload, "metadata.user_id", generateID()) + userID, errUserID := generateID() + if errUserID != nil { + return nil, errUserID + } + payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) } - return payload + return payload, nil } // fingerprintSalt is the salt used by Claude Code to compute the 3-char build fingerprint. @@ -1829,7 +1861,7 @@ IMPORTANT: this context may or may not be relevant to your tasks. You should not // applyCloaking applies cloaking transformations to the payload based on config and client. // Cloaking includes: system prompt injection, fake user ID, and sensitive word obfuscation. -func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte, model string, apiKey string) []byte { +func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte, model string, apiKey string) ([]byte, error) { clientUserAgent := getClientUserAgent(ctx) // Enable cch signing for OAuth tokens by default (not just experimental flag). oauthToken := isClaudeOAuthToken(apiKey) @@ -1862,7 +1894,7 @@ func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.A // Determine if cloaking should be applied if !helps.ShouldCloak(cloakMode, clientUserAgent) { - return payload + return payload, nil } // Skip system instructions for claude-3-5-haiku models @@ -1874,7 +1906,11 @@ func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.A } // Inject fake user ID - payload = injectFakeUserID(payload, apiKey, cacheUserID) + var errFakeUserID error + payload, errFakeUserID = injectFakeUserID(ctx, payload, apiKey, cacheUserID) + if errFakeUserID != nil { + return nil, errFakeUserID + } // Apply sensitive word obfuscation if len(sensitiveWords) > 0 { @@ -1882,7 +1918,7 @@ func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.A payload = helps.ObfuscateSensitiveWords(payload, matcher) } - return payload + return payload, nil } // ensureCacheControl injects cache_control breakpoints into the payload for optimal prompt caching. diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index c54ea598..5221aacd 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -2096,7 +2096,10 @@ func TestApplyCloaking_PreservesConfiguredStrictModeAndSensitiveWordsWhenModeOmi auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "key-123"}} payload := []byte(`{"system":"proxy rules","messages":[{"role":"user","content":[{"type":"text","text":"proxy access"}]}]}`) - out := applyCloaking(context.Background(), cfg, auth, payload, "claude-3-5-sonnet-20241022", "key-123") + out, errCloaking := applyCloaking(context.Background(), cfg, auth, payload, "claude-3-5-sonnet-20241022", "key-123") + if errCloaking != nil { + t.Fatalf("applyCloaking() error = %v", errCloaking) + } blocks := gjson.GetBytes(out, "system").Array() if len(blocks) != 3 { diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 776408fc..71b9f921 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -249,23 +249,28 @@ func (s codexReasoningReplayScope) valid() bool { } func applyCodexReasoningReplayCache(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope) { + updated, scope, _ := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + return updated, scope +} + +func applyCodexReasoningReplayCacheRequired(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope, error) { scope := codexReasoningReplayScopeFromRequest(ctx, from, req, opts, body) if !scope.valid() { - return body, scope + return body, scope, nil } - items, ok := internalcache.GetCodexReasoningReplayItems(scope.modelName, scope.sessionKey) - if !ok { - return body, scope + items, ok, errReplay := internalcache.GetCodexReasoningReplayItemsRequired(ctx, scope.modelName, scope.sessionKey) + if errReplay != nil || !ok { + return body, scope, errReplay } items = filterCodexReasoningReplayItemsForInput(body, items) if len(items) == 0 { - return body, scope + return body, scope, nil } updated, ok := insertCodexReasoningReplayItems(body, items) if !ok { - return body, scope + return body, scope, nil } - return updated, scope + return updated, scope, nil } func codexReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) codexReasoningReplayScope { @@ -299,23 +304,25 @@ func codexClaudeCodePromptCacheStorageKey(req cliproxyexecutor.Request) string { if sessionID == "" { return "" } - return fmt.Sprintf("%s-claude:%s", req.Model, sessionID) + return helps.CodexPromptCacheKey(req.Model, "claude:"+sessionID) } -func codexClaudeCodePromptCache(req cliproxyexecutor.Request) (helps.CodexCache, bool) { +func codexClaudeCodePromptCache(ctx context.Context, req cliproxyexecutor.Request) (helps.CodexCache, bool, error) { key := codexClaudeCodePromptCacheStorageKey(req) if key == "" { - return helps.CodexCache{}, false + return helps.CodexCache{}, false, nil } - if cache, ok := helps.GetCodexCache(key); ok { - return cache, true + if cache, ok, errCache := helps.GetCodexCacheRequired(ctx, key); errCache != nil || ok { + return cache, ok, errCache } cache := helps.CodexCache{ ID: uuid.New().String(), Expire: time.Now().Add(1 * time.Hour), } - helps.SetCodexCache(key, cache) - return cache, true + if errSet := helps.SetCodexCacheRequired(ctx, key, cache); errSet != nil { + return helps.CodexCache{}, false, errSet + } + return cache, true, nil } func extractClaudeCodeSessionIDForCodexReplay(payload []byte) string { @@ -724,19 +731,20 @@ func cacheCodexReasoningReplayFromCompleted(scope codexReasoningReplayScope, com continue } } - if !internalcache.CacheCodexReasoningReplayItems(scope.modelName, scope.sessionKey, items) { + if !internalcache.CacheCodexReasoningReplayItemsBestEffort(context.Background(), scope.modelName, scope.sessionKey, items) { internalcache.DeleteCodexReasoningReplayItem(scope.modelName, scope.sessionKey) } } -func clearCodexReasoningReplayOnInvalidSignature(scope codexReasoningReplayScope, statusCode int, body []byte) { +func clearCodexReasoningReplayOnInvalidSignature(ctx context.Context, scope codexReasoningReplayScope, statusCode int, body []byte) error { if !scope.valid() { - return + return nil } code, _, ok := codexStatusErrorClassification(statusCode, body) if ok && code == "thinking_signature_invalid" { - internalcache.DeleteCodexReasoningReplayItem(scope.modelName, scope.sessionKey) + return internalcache.DeleteCodexReasoningReplayItemRequired(ctx, scope.modelName, scope.sessionKey) } + return nil } // PrepareRequest injects Codex credentials into the outgoing HTTP request. @@ -818,7 +826,10 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re body = ensureImageGenerationTool(body, baseModel, auth) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) - body, replayScope := applyCodexReasoningReplayCache(ctx, from, req, opts, body) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return resp, errReplay + } reporter.SetTranslatedReasoningEffort(body, to.String()) url := strings.TrimSuffix(baseURL, "/") + "/responses" @@ -862,7 +873,9 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { b, _ := io.ReadAll(httpResp.Body) b = applyCodexIdentityConfuseResponsePayload(b, identityState) - clearCodexReasoningReplayOnInvalidSignature(replayScope, httpResp.StatusCode, b) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClearReplay != nil { + return resp, errClearReplay + } helps.AppendAPIResponseChunk(ctx, e.cfg, b) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) err = newCodexStatusErr(httpResp.StatusCode, b) @@ -888,7 +901,9 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re eventType := gjson.GetBytes(eventData, "type").String() if streamErr, terminalBody, ok := codexTerminalStreamErr(eventData); ok { - clearCodexReasoningReplayOnInvalidSignature(replayScope, streamErr.StatusCode(), terminalBody) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + return resp, errClearReplay + } err = streamErr return resp, err } @@ -1095,7 +1110,10 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au body = ensureImageGenerationTool(body, baseModel, auth) } body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) - body, replayScope := applyCodexReasoningReplayCache(ctx, from, req, opts, body) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return nil, errReplay + } reporter.SetTranslatedReasoningEffort(body, to.String()) url := strings.TrimSuffix(baseURL, "/") + "/responses" @@ -1142,7 +1160,9 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au return nil, readErr } data = applyCodexIdentityConfuseResponsePayload(data, identityState) - clearCodexReasoningReplayOnInvalidSignature(replayScope, httpResp.StatusCode, data) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, data); errClearReplay != nil { + return nil, errClearReplay + } helps.AppendAPIResponseChunk(ctx, e.cfg, data) helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) err = newCodexStatusErr(httpResp.StatusCode, data) @@ -1169,7 +1189,15 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au if bytes.HasPrefix(line, dataTag) { data := bytes.TrimSpace(line[5:]) if streamErr, terminalBody, ok := codexTerminalStreamErr(data); ok { - clearCodexReasoningReplayOnInvalidSignature(replayScope, streamErr.StatusCode(), terminalBody) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errClearReplay}: + case <-ctx.Done(): + } + return + } helps.RecordAPIResponseError(ctx, e.cfg, streamErr) reporter.PublishFailure(ctx, streamErr) select { @@ -1430,7 +1458,11 @@ type codexIdentityReplacement struct { func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Format, url string, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, userPayload []byte, rawJSON []byte) (*http.Request, []byte, codexIdentityConfuseState, error) { var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - if cached, ok := codexClaudeCodePromptCache(req); ok { + cached, ok, errCache := codexClaudeCodePromptCache(ctx, req) + if errCache != nil { + return nil, nil, codexIdentityConfuseState{}, errCache + } + if ok { cache = cached } } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 603d20e5..30ae848e 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -221,7 +221,10 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut return resp, err } - body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body) + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body) + if errPromptCache != nil { + return resp, errPromptCache + } clientBody := body var identityState codexIdentityConfuseState upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) @@ -437,7 +440,10 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr return nil, err } - body, wsHeaders := applyCodexPromptCacheHeaders(from, req, body) + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body) + if errPromptCache != nil { + return nil, errPromptCache + } clientBody := body var identityState codexIdentityConfuseState upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, userPayload, body) @@ -831,14 +837,23 @@ func buildCodexResponsesWebsocketURL(httpURL string) (string, error) { } func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header) { + body, headers, _ := applyCodexPromptCacheHeadersWithContext(context.Background(), from, req, rawJSON) + return body, headers +} + +func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header, error) { headers := http.Header{} if len(rawJSON) == 0 { - return rawJSON, headers + return rawJSON, headers, nil } var cache helps.CodexCache if sourceFormatEqual(from, sdktranslator.FormatClaude) { - if cached, ok := codexClaudeCodePromptCache(req); ok { + cached, ok, errCache := codexClaudeCodePromptCache(ctx, req) + if errCache != nil { + return nil, nil, errCache + } + if ok { cache = cached } } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { @@ -853,7 +868,7 @@ func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecuto headers.Set("Conversation_id", cache.ID) } - return rawJSON, headers + return rawJSON, headers, nil } func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config) http.Header { diff --git a/internal/runtime/executor/helps/cache_helpers.go b/internal/runtime/executor/helps/cache_helpers.go index ec063384..b52afe04 100644 --- a/internal/runtime/executor/helps/cache_helpers.go +++ b/internal/runtime/executor/helps/cache_helpers.go @@ -1,8 +1,11 @@ package helps import ( + "context" "sync" "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" ) type CodexCache struct { @@ -49,20 +52,77 @@ func purgeExpiredCodexCache() { // GetCodexCache retrieves a cached entry, returning ok=false if not found or expired. func GetCodexCache(key string) (CodexCache, bool) { + cache, ok, err := GetCodexCacheRequired(context.Background(), key) + if err == nil { + return cache, ok + } + return CodexCache{}, false +} + +// GetCodexCacheRequired retrieves a cached entry for request-time paths. +func GetCodexCacheRequired(ctx context.Context, key string) (CodexCache, bool, error) { + var homeCache CodexCache + homeMode, found, errGet := homekv.KVGetJSONRequired(ctx, key, &homeCache) + if homeMode { + if errGet != nil || !found { + return CodexCache{}, false, errGet + } + if homeCache.Expire.Before(time.Now()) { + _, _, _ = homekv.KVDelRequired(ctx, key) + return CodexCache{}, false, nil + } + return homeCache, true, nil + } + codexCacheCleanupOnce.Do(startCodexCacheCleanup) codexCacheMu.RLock() cache, ok := codexCacheMap[key] codexCacheMu.RUnlock() if !ok || cache.Expire.Before(time.Now()) { - return CodexCache{}, false + return CodexCache{}, false, nil } - return cache, true + return cache, true, nil } // SetCodexCache stores a cache entry. func SetCodexCache(key string, cache CodexCache) { + SetCodexCacheBestEffort(context.Background(), key, cache) +} + +// SetCodexCacheRequired stores a cache entry for request-time paths. +func SetCodexCacheRequired(ctx context.Context, key string, cache CodexCache) error { + ttl := time.Until(cache.Expire) + if ttl <= 0 { + return nil + } + if _, homeMode, _ := homekv.CurrentKVClient(); homeMode { + _, errSet := homekv.KVSetJSONRequired(ctx, key, cache, ttl) + return errSet + } codexCacheCleanupOnce.Do(startCodexCacheCleanup) codexCacheMu.Lock() codexCacheMap[key] = cache codexCacheMu.Unlock() + return nil +} + +// SetCodexCacheBestEffort stores a cache entry without failing completed responses. +func SetCodexCacheBestEffort(ctx context.Context, key string, cache CodexCache) bool { + ttl := time.Until(cache.Expire) + if ttl <= 0 { + return false + } + if _, homeMode, _ := homekv.CurrentKVClient(); homeMode { + return homekv.KVSetJSONBestEffort(ctx, key, cache, ttl) + } + codexCacheCleanupOnce.Do(startCodexCacheCleanup) + codexCacheMu.Lock() + codexCacheMap[key] = cache + codexCacheMu.Unlock() + return true +} + +// CodexPromptCacheKey builds the Home KV key for a model/user prompt cache. +func CodexPromptCacheKey(modelName string, userScope string) string { + return "cpa:codex:prompt-cache:" + homekv.HashKeyPart(modelName) + ":" + homekv.HashKeyPart(userScope) } diff --git a/internal/runtime/executor/helps/cache_helpers_test.go b/internal/runtime/executor/helps/cache_helpers_test.go new file mode 100644 index 00000000..3b932818 --- /dev/null +++ b/internal/runtime/executor/helps/cache_helpers_test.go @@ -0,0 +1,27 @@ +package helps + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" +) + +func TestSetCodexCacheRequiredHomeUnavailableReturnsError(t *testing.T) { + homekv.SetCurrent(homekv.New(config.HomeConfig{Enabled: false})) + t.Cleanup(homekv.ClearCurrent) + + errSet := SetCodexCacheRequired(context.Background(), "cpa:codex:prompt-cache:test", CodexCache{ + ID: "cache-id", + Expire: time.Now().Add(time.Hour), + }) + if errSet == nil { + t.Fatal("SetCodexCacheRequired() error = nil, want home kv unavailable error") + } + if !strings.Contains(errSet.Error(), "home kv store unavailable") { + t.Fatalf("SetCodexCacheRequired() error = %v, want home kv store unavailable", errSet) + } +} diff --git a/internal/runtime/executor/helps/claude_device_profile.go b/internal/runtime/executor/helps/claude_device_profile.go index 09f04929..2eb97d98 100644 --- a/internal/runtime/executor/helps/claude_device_profile.go +++ b/internal/runtime/executor/helps/claude_device_profile.go @@ -1,8 +1,11 @@ package helps import ( + "context" "crypto/sha256" "encoding/hex" + "encoding/json" + "fmt" "net/http" "regexp" "runtime" @@ -12,6 +15,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) @@ -22,6 +26,7 @@ const ( defaultClaudeFingerprintOS = "MacOS" defaultClaudeFingerprintArch = "arm64" claudeDeviceProfileTTL = 7 * 24 * time.Hour + claudeDeviceProfileLockTTL = 5 * time.Second claudeDeviceProfileCleanupPeriod = time.Hour ) @@ -35,6 +40,17 @@ var ( ClaudeDeviceProfileBeforeCandidateStore func(ClaudeDeviceProfile) ) +type claudeDeviceProfileKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (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 currentClaudeDeviceProfileKVClient = func() (claudeDeviceProfileKVClient, bool, error) { + return homekv.CurrentKVClient() +} + type claudeCLIVersion struct { major int minor int @@ -78,6 +94,14 @@ type claudeDeviceProfileCacheEntry struct { expire time.Time } +type claudeDeviceProfileKVValue struct { + UserAgent string `json:"user_agent"` + PackageVersion string `json:"package_version"` + RuntimeVersion string `json:"runtime_version"` + OS string `json:"os"` + Arch string `json:"arch"` +} + func ClaudeDeviceProfileStabilizationEnabled(cfg *config.Config) bool { if cfg == nil || cfg.ClaudeHeaderDefaults.StabilizeDeviceProfile == nil { return false @@ -256,6 +280,14 @@ func claudeDeviceProfileCacheKey(auth *cliproxyauth.Auth, apiKey string) string return hex.EncodeToString(sum[:]) } +func claudeDeviceProfileKVKey(auth *cliproxyauth.Auth, apiKey string) string { + return "cpa:claude:device-profile:" + homekv.HashKeyPart(claudeDeviceProfileScopeKey(auth, apiKey)) +} + +func claudeDeviceProfileLockKVKey(auth *cliproxyauth.Auth, apiKey string) string { + return "cpa:claude:device-profile-lock:" + homekv.HashKeyPart(claudeDeviceProfileScopeKey(auth, apiKey)) +} + func startClaudeDeviceProfileCacheCleanup() { go func() { ticker := time.NewTicker(claudeDeviceProfileCleanupPeriod) @@ -278,6 +310,26 @@ func purgeExpiredClaudeDeviceProfiles() { } func ResolveClaudeDeviceProfile(auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) ClaudeDeviceProfile { + profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, apiKey, headers, cfg) + if errProfile != nil { + return defaultClaudeDeviceProfile(cfg) + } + return profile +} + +// ResolveClaudeDeviceProfileRequired resolves a stable Claude Code device profile for request-time paths. +func ResolveClaudeDeviceProfileRequired(ctx context.Context, auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) (ClaudeDeviceProfile, error) { + client, homeMode, errClient := currentClaudeDeviceProfileKVClient() + if homeMode { + if errClient != nil { + return ClaudeDeviceProfile{}, errClient + } + return resolveClaudeDeviceProfileHome(ctx, client, auth, apiKey, headers, cfg) + } + return resolveClaudeDeviceProfileLocal(auth, apiKey, headers, cfg), nil +} + +func resolveClaudeDeviceProfileLocal(auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) ClaudeDeviceProfile { claudeDeviceProfileCacheCleanupOnce.Do(startClaudeDeviceProfileCacheCleanup) cacheKey := claudeDeviceProfileCacheKey(auth, apiKey) @@ -338,6 +390,123 @@ func ResolveClaudeDeviceProfile(auth *cliproxyauth.Auth, apiKey string, headers return baseline } +func resolveClaudeDeviceProfileHome(ctx context.Context, client claudeDeviceProfileKVClient, auth *cliproxyauth.Auth, apiKey string, headers http.Header, cfg *config.Config) (ClaudeDeviceProfile, error) { + baseline := defaultClaudeDeviceProfile(cfg) + candidate, hasCandidate := extractClaudeDeviceProfile(headers, cfg) + if hasCandidate { + candidate = pinClaudeDeviceProfilePlatform(candidate, baseline) + } + if hasCandidate && !shouldUpgradeClaudeDeviceProfile(candidate, baseline) { + hasCandidate = false + } + + valueKey := claudeDeviceProfileKVKey(auth, apiKey) + if !hasCandidate { + return readClaudeDeviceProfileFromHome(ctx, client, valueKey, baseline) + } + + lockKey := claudeDeviceProfileLockKVKey(auth, apiKey) + gotLock, errLock := client.KVSetNX(ctx, lockKey, []byte("1"), claudeDeviceProfileLockTTL) + if errLock != nil { + return ClaudeDeviceProfile{}, errLock + } + if ClaudeDeviceProfileBeforeCandidateStore != nil { + ClaudeDeviceProfileBeforeCandidateStore(candidate) + } + + cached, found, errRead := readClaudeDeviceProfileValueFromHome(ctx, client, valueKey, baseline) + if errRead != nil { + return ClaudeDeviceProfile{}, errRead + } + if found && !shouldUpgradeClaudeDeviceProfile(candidate, cached) { + if _, errExpire := client.KVExpire(ctx, valueKey, claudeDeviceProfileTTL); errExpire != nil { + return ClaudeDeviceProfile{}, errExpire + } + return cached, nil + } + if !gotLock { + if found { + return cached, nil + } + return ClaudeDeviceProfile{}, fmt.Errorf("home kv device profile lock not acquired and profile missing") + } + + if errWrite := writeClaudeDeviceProfileToHome(ctx, client, valueKey, candidate); errWrite != nil { + return ClaudeDeviceProfile{}, errWrite + } + return candidate, nil +} + +func readClaudeDeviceProfileFromHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, baseline ClaudeDeviceProfile) (ClaudeDeviceProfile, error) { + profile, found, errRead := readClaudeDeviceProfileValueFromHome(ctx, client, key, baseline) + if errRead != nil { + return ClaudeDeviceProfile{}, errRead + } + if !found { + return baseline, nil + } + if _, errExpire := client.KVExpire(ctx, key, claudeDeviceProfileTTL); errExpire != nil { + return ClaudeDeviceProfile{}, errExpire + } + return profile, nil +} + +func readClaudeDeviceProfileValueFromHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, baseline ClaudeDeviceProfile) (ClaudeDeviceProfile, bool, error) { + raw, found, errGet := client.KVGet(ctx, key) + if errGet != nil || !found { + return ClaudeDeviceProfile{}, false, errGet + } + var value claudeDeviceProfileKVValue + if errUnmarshal := json.Unmarshal(raw, &value); errUnmarshal != nil { + return ClaudeDeviceProfile{}, false, errUnmarshal + } + profile := value.ToProfile() + if strings.TrimSpace(profile.UserAgent) == "" { + return ClaudeDeviceProfile{}, false, nil + } + return normalizeClaudeDeviceProfile(profile, baseline), true, nil +} + +func writeClaudeDeviceProfileToHome(ctx context.Context, client claudeDeviceProfileKVClient, key string, profile ClaudeDeviceProfile) error { + raw, errMarshal := json.Marshal(claudeDeviceProfileKVValueFromProfile(profile)) + if errMarshal != nil { + return errMarshal + } + written, errSet := client.KVSet(ctx, key, raw, homekv.KVSetOptions{EX: claudeDeviceProfileTTL}) + if errSet != nil { + return errSet + } + if !written { + return fmt.Errorf("home kv device profile write skipped") + } + return nil +} + +func claudeDeviceProfileKVValueFromProfile(profile ClaudeDeviceProfile) claudeDeviceProfileKVValue { + return claudeDeviceProfileKVValue{ + UserAgent: profile.UserAgent, + PackageVersion: profile.PackageVersion, + RuntimeVersion: profile.RuntimeVersion, + OS: profile.OS, + Arch: profile.Arch, + } +} + +func (value claudeDeviceProfileKVValue) ToProfile() ClaudeDeviceProfile { + profile := ClaudeDeviceProfile{ + UserAgent: strings.TrimSpace(value.UserAgent), + PackageVersion: strings.TrimSpace(value.PackageVersion), + RuntimeVersion: strings.TrimSpace(value.RuntimeVersion), + OS: strings.TrimSpace(value.OS), + Arch: strings.TrimSpace(value.Arch), + } + if version, ok := parseClaudeCLIVersion(profile.UserAgent); ok { + profile.version = version + profile.hasVersion = true + } + return profile +} + func ApplyClaudeDeviceProfileHeaders(r *http.Request, profile ClaudeDeviceProfile) { if r == nil { return diff --git a/internal/runtime/executor/helps/claude_device_profile_test.go b/internal/runtime/executor/helps/claude_device_profile_test.go new file mode 100644 index 00000000..0f99168d --- /dev/null +++ b/internal/runtime/executor/helps/claude_device_profile_test.go @@ -0,0 +1,237 @@ +package helps + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +type fakeClaudeDeviceProfileKVClient struct { + values map[string][]byte + getErr error + setErr error + setNXErr error + expireErr error + setNXResult bool + getCount int + setCount int + setNXCount int + expireCount int + lastSetTTL time.Duration + lastSetNXTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeClaudeDeviceProfileKVClient() *fakeClaudeDeviceProfileKVClient { + return &fakeClaudeDeviceProfileKVClient{ + values: make(map[string][]byte), + setNXResult: true, + } +} + +func (c *fakeClaudeDeviceProfileKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeClaudeDeviceProfileKVClient) KVSet(_ context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) { + c.setCount++ + c.lastSetTTL = opts.EX + if c.setErr != nil { + return false, c.setErr + } + c.values[key] = append([]byte(nil), value...) + return true, nil +} + +func (c *fakeClaudeDeviceProfileKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) { + c.setNXCount++ + c.lastSetNXTTL = ttl + if c.setNXErr != nil { + return false, c.setNXErr + } + if _, ok := c.values[key]; ok { + return false, nil + } + if c.setNXResult { + c.values[key] = append([]byte(nil), value...) + return true, nil + } + return false, nil +} + +func (c *fakeClaudeDeviceProfileKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.expireCount++ + c.lastExpireTTL = ttl + if c.expireErr != nil { + return false, c.expireErr + } + return true, nil +} + +func useFakeClaudeDeviceProfileKVClient(t *testing.T, client *fakeClaudeDeviceProfileKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentClaudeDeviceProfileKVClient + currentClaudeDeviceProfileKVClient = func() (claudeDeviceProfileKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentClaudeDeviceProfileKVClient = previous + }) +} + +func mustClaudeDeviceProfileJSON(t *testing.T, value claudeDeviceProfileKVValue) []byte { + t.Helper() + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + t.Fatalf("marshal device profile: %v", errMarshal) + } + return raw +} + +func claudeDeviceHeaders(userAgent string) http.Header { + return http.Header{ + "User-Agent": {userAgent}, + "X-Stainless-Package-Version": {"0.80.0"}, + "X-Stainless-Runtime-Version": {"v24.4.0"}, + "X-Stainless-Os": {"Windows"}, + "X-Stainless-Arch": {"x64"}, + } +} + +func TestResolveClaudeDeviceProfileRequiredHomeReadWithoutCandidate(t *testing.T) { + client := newFakeClaudeDeviceProfileKVClient() + auth := &cliproxyauth.Auth{ID: "auth-1"} + key := claudeDeviceProfileKVKey(auth, "api-key") + client.values[key] = mustClaudeDeviceProfileJSON(t, claudeDeviceProfileKVValue{ + UserAgent: "claude-cli/2.2.0 (external, cli)", + PackageVersion: "0.80.0", + RuntimeVersion: "v24.4.0", + OS: "Windows", + Arch: "x64", + }) + useFakeClaudeDeviceProfileKVClient(t, client, true, nil) + + profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", nil, nil) + if errProfile != nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile) + } + if profile.UserAgent != "claude-cli/2.2.0 (external, cli)" { + t.Fatalf("UserAgent = %q, want cached profile", profile.UserAgent) + } + if profile.OS != defaultClaudeFingerprintOS || profile.Arch != defaultClaudeFingerprintArch { + t.Fatalf("platform = %s/%s, want baseline pinned %s/%s", profile.OS, profile.Arch, defaultClaudeFingerprintOS, defaultClaudeFingerprintArch) + } + if client.expireCount != 1 || client.lastExpireTTL != claudeDeviceProfileTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, claudeDeviceProfileTTL) + } +} + +func TestResolveClaudeDeviceProfileRequiredHomeCandidateLocksRereadsAndWrites(t *testing.T) { + client := newFakeClaudeDeviceProfileKVClient() + auth := &cliproxyauth.Auth{ID: "auth-1"} + useFakeClaudeDeviceProfileKVClient(t, client, true, nil) + + profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), nil) + if errProfile != nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile) + } + if profile.UserAgent != "claude-cli/2.2.0 (external, cli)" { + t.Fatalf("UserAgent = %q, want candidate", profile.UserAgent) + } + if client.setNXCount != 1 || client.lastSetNXTTL != claudeDeviceProfileLockTTL { + t.Fatalf("KVSetNX count/ttl = %d/%v, want 1/%v", client.setNXCount, client.lastSetNXTTL, claudeDeviceProfileLockTTL) + } + if client.getCount != 1 { + t.Fatalf("KVGet count = %d, want re-read after lock", client.getCount) + } + if client.setCount != 1 || client.lastSetTTL != claudeDeviceProfileTTL { + t.Fatalf("KVSet count/ttl = %d/%v, want 1/%v", client.setCount, client.lastSetTTL, claudeDeviceProfileTTL) + } +} + +func TestResolveClaudeDeviceProfileRequiredHomeCandidateDoesNotDowngradeCachedProfile(t *testing.T) { + client := newFakeClaudeDeviceProfileKVClient() + auth := &cliproxyauth.Auth{ID: "auth-1"} + key := claudeDeviceProfileKVKey(auth, "api-key") + client.values[key] = mustClaudeDeviceProfileJSON(t, claudeDeviceProfileKVValue{ + UserAgent: "claude-cli/2.4.0 (external, cli)", + PackageVersion: "0.90.0", + RuntimeVersion: "v24.5.0", + OS: "Windows", + Arch: "x64", + }) + useFakeClaudeDeviceProfileKVClient(t, client, true, nil) + + profile, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders("claude-cli/2.3.0 (external, cli)"), nil) + if errProfile != nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() error = %v", errProfile) + } + if profile.UserAgent != "claude-cli/2.4.0 (external, cli)" { + t.Fatalf("UserAgent = %q, want higher cached profile", profile.UserAgent) + } + if client.setCount != 0 { + t.Fatalf("KVSet count = %d, want no downgrade write", client.setCount) + } + if client.expireCount != 1 { + t.Fatalf("KVExpire count = %d, want cached refresh", client.expireCount) + } +} + +func TestResolveClaudeDeviceProfileRequiredHomeFailures(t *testing.T) { + for _, tc := range []struct { + name string + headers http.Header + client *fakeClaudeDeviceProfileKVClient + }{ + {name: "read", client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "lock", headers: claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, setNXErr: errors.New("lock failed")}}, + {name: "lock-miss", headers: claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: false}}, + {name: "reread", headers: claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, getErr: errors.New("re-read failed")}}, + {name: "write", headers: claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), client: &fakeClaudeDeviceProfileKVClient{values: make(map[string][]byte), setNXResult: true, setErr: errors.New("write failed")}}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeClaudeDeviceProfileKVClient(t, tc.client, true, nil) + if _, errProfile := ResolveClaudeDeviceProfileRequired(context.Background(), &cliproxyauth.Auth{ID: "auth-1"}, "api-key", tc.headers, nil); errProfile == nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() error = nil, want error") + } + }) + } +} + +func TestResolveClaudeDeviceProfileRequiredNonHomeKeepsLocalCache(t *testing.T) { + ResetClaudeDeviceProfileCache() + client := newFakeClaudeDeviceProfileKVClient() + useFakeClaudeDeviceProfileKVClient(t, client, false, nil) + auth := &cliproxyauth.Auth{ID: "auth-1"} + cfg := &config.Config{} + + first, errFirst := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", claudeDeviceHeaders("claude-cli/2.2.0 (external, cli)"), cfg) + if errFirst != nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() first error = %v", errFirst) + } + second, errSecond := ResolveClaudeDeviceProfileRequired(context.Background(), auth, "api-key", nil, cfg) + if errSecond != nil { + t.Fatalf("ResolveClaudeDeviceProfileRequired() second error = %v", errSecond) + } + if second.UserAgent != first.UserAgent { + t.Fatalf("cached UserAgent = %q, want %q", second.UserAgent, first.UserAgent) + } + if client.getCount != 0 || client.setCount != 0 || client.setNXCount != 0 { + t.Fatalf("KV calls = get %d set %d setnx %d, want all zero", client.getCount, client.setCount, client.setNXCount) + } +} diff --git a/internal/runtime/executor/helps/session_id_cache.go b/internal/runtime/executor/helps/session_id_cache.go index 6c89f001..015fb3e3 100644 --- a/internal/runtime/executor/helps/session_id_cache.go +++ b/internal/runtime/executor/helps/session_id_cache.go @@ -1,12 +1,16 @@ package 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 { @@ -20,6 +24,16 @@ var ( 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 @@ -53,8 +67,46 @@ func sessionIDCacheKey(apiKey string) string { // 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() + 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) @@ -73,7 +125,7 @@ func CachedSessionID(apiKey string) string { entry.expire = now.Add(sessionIDTTL) sessionIDCache[key] = entry sessionIDCacheMu.Unlock() - return entry.value + return entry.value, nil } sessionIDCacheMu.Unlock() } @@ -88,5 +140,9 @@ func CachedSessionID(apiKey string) string { entry.expire = now.Add(sessionIDTTL) sessionIDCache[key] = entry sessionIDCacheMu.Unlock() - return entry.value + return entry.value, nil +} + +func claudeSessionIDKVKey(apiKey string) string { + return "cpa:claude:session-id:" + homekv.HashKeyPart(apiKey) } diff --git a/internal/runtime/executor/helps/session_id_cache_test.go b/internal/runtime/executor/helps/session_id_cache_test.go new file mode 100644 index 00000000..ef890666 --- /dev/null +++ b/internal/runtime/executor/helps/session_id_cache_test.go @@ -0,0 +1,178 @@ +package helps + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" +) + +func resetSessionIDCache() { + sessionIDCacheMu.Lock() + sessionIDCache = make(map[string]sessionIDCacheEntry) + sessionIDCacheMu.Unlock() +} + +type fakeClaudeIDKVClient struct { + values map[string][]byte + getErr error + setErr error + expireErr error + setNoPersist bool + getCount int + setCount int + expireCount int + lastSetTTL time.Duration + lastExpireTTL time.Duration +} + +func newFakeClaudeIDKVClient() *fakeClaudeIDKVClient { + return &fakeClaudeIDKVClient{values: make(map[string][]byte)} +} + +func (c *fakeClaudeIDKVClient) KVGet(_ context.Context, key string) ([]byte, bool, error) { + c.getCount++ + if c.getErr != nil { + return nil, false, c.getErr + } + value, ok := c.values[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), value...), true, nil +} + +func (c *fakeClaudeIDKVClient) KVSetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) { + c.setCount++ + c.lastSetTTL = ttl + if c.setErr != nil { + return false, c.setErr + } + if _, ok := c.values[key]; ok { + return false, nil + } + if !c.setNoPersist { + c.values[key] = append([]byte(nil), value...) + } + return true, nil +} + +func (c *fakeClaudeIDKVClient) KVExpire(_ context.Context, _ string, ttl time.Duration) (bool, error) { + c.expireCount++ + c.lastExpireTTL = ttl + if c.expireErr != nil { + return false, c.expireErr + } + return true, nil +} + +func useFakeClaudeIDKVClient(t *testing.T, client *fakeClaudeIDKVClient, homeMode bool, errClient error) { + t.Helper() + previous := currentClaudeIDKVClient + currentClaudeIDKVClient = func() (claudeIDKVClient, bool, error) { + return client, homeMode, errClient + } + t.Cleanup(func() { + currentClaudeIDKVClient = previous + }) +} + +func TestCachedSessionIDRequiredHomeReusesKVAcrossLocalCacheReset(t *testing.T) { + resetSessionIDCache() + client := newFakeClaudeIDKVClient() + useFakeClaudeIDKVClient(t, client, true, nil) + + first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1") + if errFirst != nil { + t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst) + } + resetSessionIDCache() + second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1") + if errSecond != nil { + t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond) + } + if first != second { + t.Fatalf("session id = %q then %q, want same Home KV value", first, second) + } + if _, errParse := uuid.Parse(first); errParse != nil { + t.Fatalf("session id %q is not a UUID: %v", first, errParse) + } + if client.setCount != 1 { + t.Fatalf("KVSetNX count = %d, want 1", client.setCount) + } + if client.expireCount != 1 || client.lastExpireTTL != sessionIDTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, sessionIDTTL) + } + if client.lastSetTTL != sessionIDTTL { + t.Fatalf("KVSetNX ttl = %v, want %v", client.lastSetTTL, sessionIDTTL) + } +} + +func TestCachedSessionIDRequiredEmptyAPIKeyDoesNotUseHomeKV(t *testing.T) { + client := newFakeClaudeIDKVClient() + useFakeClaudeIDKVClient(t, client, true, nil) + + value, errValue := CachedSessionIDRequired(context.Background(), "") + if errValue != nil { + t.Fatalf("CachedSessionIDRequired(empty) error = %v", errValue) + } + if _, errParse := uuid.Parse(value); errParse != nil { + t.Fatalf("session id %q is not a UUID: %v", value, errParse) + } + if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 { + t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount) + } +} + +func TestCachedSessionIDRequiredHomeKVFailures(t *testing.T) { + for _, tc := range []struct { + name string + client *fakeClaudeIDKVClient + }{ + {name: "get", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "set", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}}, + {name: "expire", client: &fakeClaudeIDKVClient{values: map[string][]byte{ + claudeSessionIDKVKey("api-key-1"): []byte(uuid.New().String()), + }, expireErr: errors.New("expire failed")}}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeClaudeIDKVClient(t, tc.client, true, nil) + if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil { + t.Fatalf("CachedSessionIDRequired() error = nil, want error") + } + }) + } +} + +func TestCachedSessionIDRequiredHomeRequiresReadAfterSet(t *testing.T) { + client := newFakeClaudeIDKVClient() + client.setNoPersist = true + useFakeClaudeIDKVClient(t, client, true, nil) + + if _, errValue := CachedSessionIDRequired(context.Background(), "api-key-1"); errValue == nil { + t.Fatalf("CachedSessionIDRequired() error = nil, want missing-after-set error") + } +} + +func TestCachedSessionIDRequiredNonHomeModeUsesLocalMap(t *testing.T) { + resetSessionIDCache() + client := newFakeClaudeIDKVClient() + useFakeClaudeIDKVClient(t, client, false, nil) + + first, errFirst := CachedSessionIDRequired(context.Background(), "api-key-1") + if errFirst != nil { + t.Fatalf("CachedSessionIDRequired() first error = %v", errFirst) + } + second, errSecond := CachedSessionIDRequired(context.Background(), "api-key-1") + if errSecond != nil { + t.Fatalf("CachedSessionIDRequired() second error = %v", errSecond) + } + if first != second { + t.Fatalf("session id = %q then %q, want local cache reuse", first, second) + } + if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 { + t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount) + } +} diff --git a/internal/runtime/executor/helps/user_id_cache.go b/internal/runtime/executor/helps/user_id_cache.go index ad41fd9a..7ed87132 100644 --- a/internal/runtime/executor/helps/user_id_cache.go +++ b/internal/runtime/executor/helps/user_id_cache.go @@ -1,10 +1,15 @@ package helps import ( + "context" "crypto/sha256" "encoding/hex" + "fmt" + "strings" "sync" "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" ) type userIDCacheEntry struct { @@ -50,8 +55,46 @@ func userIDCacheKey(apiKey string) string { } func CachedUserID(apiKey string) string { + value, errValue := CachedUserIDRequired(context.Background(), apiKey) + if errValue == nil && value != "" { + return value + } + return generateFakeUserID() +} + +// CachedUserIDRequired returns a stable fake user ID per apiKey for request-time paths. +func CachedUserIDRequired(ctx context.Context, apiKey string) (string, error) { if apiKey == "" { - return generateFakeUserID() + return generateFakeUserID(), nil + } + client, homeMode, errClient := currentClaudeIDKVClient() + if homeMode { + if errClient != nil { + return "", errClient + } + key := claudeUserIDKVKey(apiKey) + raw, found, errGet := client.KVGet(ctx, key) + if errGet != nil { + return "", errGet + } + if found && isValidUserID(strings.TrimSpace(string(raw))) { + if _, errExpire := client.KVExpire(ctx, key, userIDTTL); errExpire != nil { + return "", errExpire + } + return strings.TrimSpace(string(raw)), nil + } + newID := generateFakeUserID() + if _, errSet := client.KVSetNX(ctx, key, []byte(newID), userIDTTL); errSet != nil { + return "", errSet + } + raw, found, errGet = client.KVGet(ctx, key) + if errGet != nil { + return "", errGet + } + if found && isValidUserID(strings.TrimSpace(string(raw))) { + return strings.TrimSpace(string(raw)), nil + } + return "", fmt.Errorf("home kv user id missing after set") } userIDCacheCleanupOnce.Do(startUserIDCacheCleanup) @@ -70,7 +113,7 @@ func CachedUserID(apiKey string) string { entry.expire = now.Add(userIDTTL) userIDCache[key] = entry userIDCacheMu.Unlock() - return entry.value + return entry.value, nil } userIDCacheMu.Unlock() } @@ -85,5 +128,9 @@ func CachedUserID(apiKey string) string { entry.expire = now.Add(userIDTTL) userIDCache[key] = entry userIDCacheMu.Unlock() - return entry.value + return entry.value, nil +} + +func claudeUserIDKVKey(apiKey string) string { + return "cpa:claude:user-id:" + homekv.HashKeyPart(apiKey) } diff --git a/internal/runtime/executor/helps/user_id_cache_test.go b/internal/runtime/executor/helps/user_id_cache_test.go index b166576c..ed0a663c 100644 --- a/internal/runtime/executor/helps/user_id_cache_test.go +++ b/internal/runtime/executor/helps/user_id_cache_test.go @@ -1,6 +1,8 @@ package helps import ( + "context" + "errors" "testing" "time" ) @@ -84,3 +86,80 @@ func TestCachedUserID_RenewsTTLOnHit(t *testing.T) { t.Fatalf("expected TTL to renew, got %v remaining", entry.expire.Sub(soon)) } } + +func TestCachedUserIDRequiredHomeReusesKVAcrossLocalCacheReset(t *testing.T) { + resetUserIDCache() + client := newFakeClaudeIDKVClient() + useFakeClaudeIDKVClient(t, client, true, nil) + + first, errFirst := CachedUserIDRequired(context.Background(), "api-key-1") + if errFirst != nil { + t.Fatalf("CachedUserIDRequired() first error = %v", errFirst) + } + resetUserIDCache() + second, errSecond := CachedUserIDRequired(context.Background(), "api-key-1") + if errSecond != nil { + t.Fatalf("CachedUserIDRequired() second error = %v", errSecond) + } + if first != second { + t.Fatalf("user id = %q then %q, want same Home KV value", first, second) + } + if !IsValidUserID(first) { + t.Fatalf("user id %q is not valid", first) + } + if client.setCount != 1 { + t.Fatalf("KVSetNX count = %d, want 1", client.setCount) + } + if client.expireCount != 1 || client.lastExpireTTL != userIDTTL { + t.Fatalf("KVExpire count/ttl = %d/%v, want 1/%v", client.expireCount, client.lastExpireTTL, userIDTTL) + } + if client.lastSetTTL != userIDTTL { + t.Fatalf("KVSetNX ttl = %v, want %v", client.lastSetTTL, userIDTTL) + } +} + +func TestCachedUserIDRequiredEmptyAPIKeyDoesNotUseHomeKV(t *testing.T) { + client := newFakeClaudeIDKVClient() + useFakeClaudeIDKVClient(t, client, true, nil) + + value, errValue := CachedUserIDRequired(context.Background(), "") + if errValue != nil { + t.Fatalf("CachedUserIDRequired(empty) error = %v", errValue) + } + if !IsValidUserID(value) { + t.Fatalf("user id %q is not valid", value) + } + if client.getCount != 0 || client.setCount != 0 || client.expireCount != 0 { + t.Fatalf("KV calls = get %d set %d expire %d, want all zero", client.getCount, client.setCount, client.expireCount) + } +} + +func TestCachedUserIDRequiredHomeKVFailures(t *testing.T) { + for _, tc := range []struct { + name string + client *fakeClaudeIDKVClient + }{ + {name: "get", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), getErr: errors.New("get failed")}}, + {name: "set", client: &fakeClaudeIDKVClient{values: make(map[string][]byte), setErr: errors.New("set failed")}}, + {name: "expire", client: &fakeClaudeIDKVClient{values: map[string][]byte{ + claudeUserIDKVKey("api-key-1"): []byte(GenerateFakeUserID()), + }, expireErr: errors.New("expire failed")}}, + } { + t.Run(tc.name, func(t *testing.T) { + useFakeClaudeIDKVClient(t, tc.client, true, nil) + if _, errValue := CachedUserIDRequired(context.Background(), "api-key-1"); errValue == nil { + t.Fatalf("CachedUserIDRequired() error = nil, want error") + } + }) + } +} + +func TestCachedUserIDRequiredHomeRequiresReadAfterSet(t *testing.T) { + client := newFakeClaudeIDKVClient() + client.setNoPersist = true + useFakeClaudeIDKVClient(t, client, true, nil) + + if _, errValue := CachedUserIDRequired(context.Background(), "api-key-1"); errValue == nil { + t.Fatalf("CachedUserIDRequired() error = nil, want missing-after-set error") + } +} diff --git a/internal/translator/antigravity/claude/antigravity_claude_request.go b/internal/translator/antigravity/claude/antigravity_claude_request.go index d4490bc3..d196de7c 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_request.go +++ b/internal/translator/antigravity/claude/antigravity_claude_request.go @@ -6,6 +6,7 @@ package claude import ( + "context" "strings" "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" @@ -19,36 +20,56 @@ import ( ) func resolveThinkingSignature(modelName, thinkingText, rawSignature string) string { + signature, errSignature := resolveThinkingSignatureRequired(context.Background(), modelName, thinkingText, rawSignature) + if errSignature != nil { + return "" + } + return signature +} + +func resolveThinkingSignatureRequired(ctx context.Context, modelName, thinkingText, rawSignature string) (string, error) { targetProvider := sigcompat.SignatureProviderFromModelName(modelName) if targetProvider == sigcompat.SignatureProviderGemini { - return resolveProviderCompatibleSignature(targetProvider, rawSignature, sigcompat.SignatureBlockKindGeminiModelPart) + return resolveProviderCompatibleSignature(targetProvider, rawSignature, sigcompat.SignatureBlockKindGeminiModelPart), nil } if cache.SignatureCacheEnabled() { - return resolveCacheModeSignature(modelName, thinkingText, rawSignature) + return resolveCacheModeSignatureRequired(ctx, modelName, thinkingText, rawSignature) } if signature := resolveProviderCompatibleSignature(targetProvider, rawSignature, sigcompat.SignatureBlockKindUnknown); signature != "" { - return signature + return signature, nil } - return resolveBypassModeSignatureForProvider(targetProvider, rawSignature) + return resolveBypassModeSignatureForProvider(targetProvider, rawSignature), nil } func resolveCacheModeSignature(modelName, thinkingText, rawSignature string) string { + signature, errSignature := resolveCacheModeSignatureRequired(context.Background(), modelName, thinkingText, rawSignature) + if errSignature != nil { + return "" + } + return signature +} + +func resolveCacheModeSignatureRequired(ctx context.Context, modelName, thinkingText, rawSignature string) (string, error) { targetProvider := sigcompat.SignatureProviderFromModelName(modelName) if thinkingText != "" { - if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" { + cachedSig, errCachedSig := cache.GetCachedSignatureRequired(ctx, modelName, thinkingText) + if errCachedSig != nil { + return "", errCachedSig + } + if cachedSig != "" { if targetProvider == sigcompat.SignatureProviderClaude { signature, ok := sigcompat.CompatibleAntigravityClaudeThinkingSignature(cachedSig) if !ok { - return "" + return "", nil } - return signature + return signature, nil } - return cachedSig + return cachedSig, nil } } if rawSignature == "" { - return "" + return "", nil } clientSignature := "" @@ -62,14 +83,46 @@ func resolveCacheModeSignature(modelName, thinkingText, rawSignature string) str if targetProvider == sigcompat.SignatureProviderClaude { signature, ok := sigcompat.CompatibleAntigravityClaudeThinkingSignature(clientSignature) if !ok { - return "" + return "", nil } - return signature + return signature, nil } - return clientSignature + return clientSignature, nil } - return "" + return "", nil +} + +func RequireCachedThinkingSignatures(ctx context.Context, modelName string, rawJSON []byte) error { + if !cache.SignatureCacheEnabled() { + return nil + } + if sigcompat.SignatureProviderFromModelName(modelName) == sigcompat.SignatureProviderGemini { + return nil + } + messagesResult := gjson.GetBytes(rawJSON, "messages") + if !messagesResult.IsArray() { + return nil + } + for _, messageResult := range messagesResult.Array() { + contentsResult := messageResult.Get("content") + if !contentsResult.IsArray() { + continue + } + for _, contentResult := range contentsResult.Array() { + if contentResult.Get("type").String() != "thinking" { + continue + } + thinkingText := thinking.GetThinkingText(contentResult) + if thinkingText == "" { + continue + } + if _, errSignature := cache.GetCachedSignatureRequired(ctx, modelName, thinkingText); errSignature != nil { + return errSignature + } + } + } + return nil } func resolveBypassModeSignature(rawSignature string) string { diff --git a/internal/translator/antigravity/claude/antigravity_claude_response.go b/internal/translator/antigravity/claude/antigravity_claude_response.go index c883f182..da5098df 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_response.go +++ b/internal/translator/antigravity/claude/antigravity_claude_response.go @@ -100,7 +100,7 @@ var toolUseIDCounter uint64 // // Returns: // - [][]byte: A slice of bytes, each containing a Claude Code-compatible SSE payload. -func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) [][]byte { +func ConvertAntigravityResponseToClaude(ctx context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) [][]byte { if *param == nil { *param = &Params{ HasFirstResponse: false, @@ -134,7 +134,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq return } if params.CurrentThinkingText.Len() > 0 { - cache.CacheSignature(modelName, params.CurrentThinkingText.String(), signature) + cache.CacheSignatureBestEffort(ctx, modelName, params.CurrentThinkingText.String(), signature) params.CurrentThinkingText.Reset() } sigValue := formatClaudeSignatureValue(modelName, signature) diff --git a/sdk/cliproxy/auth/antigravity_credits.go b/sdk/cliproxy/auth/antigravity_credits.go index 77b03bfd..6b9480b6 100644 --- a/sdk/cliproxy/auth/antigravity_credits.go +++ b/sdk/cliproxy/auth/antigravity_credits.go @@ -5,6 +5,8 @@ import ( "strings" "sync" "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" ) type antigravityUseCreditsContextKey struct{} @@ -45,25 +47,43 @@ func SetAntigravityCreditsHint(authID string, hint AntigravityCreditsHint) { if hint.UpdatedAt.IsZero() { hint.UpdatedAt = time.Now() } + if _, homeMode, _ := homekv.CurrentKVClient(); homeMode { + homekv.KVSetJSONBestEffort(context.Background(), antigravityCreditsHintKey(authID), hint, 30*time.Minute) + return + } antigravityCreditsHintByAuth.Store(authID, hint) } // GetAntigravityCreditsHint returns the latest known AI credits state for an auth. func GetAntigravityCreditsHint(authID string) (AntigravityCreditsHint, bool) { + hint, ok, err := GetAntigravityCreditsHintRequired(context.Background(), authID) + if err == nil { + return hint, ok + } + return AntigravityCreditsHint{}, false +} + +// GetAntigravityCreditsHintRequired returns the latest known AI credits state for request-time paths. +func GetAntigravityCreditsHintRequired(ctx context.Context, authID string) (AntigravityCreditsHint, bool, error) { authID = strings.TrimSpace(authID) if authID == "" { - return AntigravityCreditsHint{}, false + return AntigravityCreditsHint{}, false, nil + } + var homeHint AntigravityCreditsHint + homeMode, found, errGet := homekv.KVGetJSONRequired(ctx, antigravityCreditsHintKey(authID), &homeHint) + if homeMode { + return homeHint, found, errGet } value, ok := antigravityCreditsHintByAuth.Load(authID) if !ok { - return AntigravityCreditsHint{}, false + return AntigravityCreditsHint{}, false, nil } hint, ok := value.(AntigravityCreditsHint) if !ok { antigravityCreditsHintByAuth.Delete(authID) - return AntigravityCreditsHint{}, false + return AntigravityCreditsHint{}, false, nil } - return hint, true + return hint, true, nil } // HasKnownAntigravityCreditsHint reports whether credits state has been discovered for an auth. @@ -72,6 +92,10 @@ func HasKnownAntigravityCreditsHint(authID string) bool { return ok && hint.Known } +func antigravityCreditsHintKey(authID string) string { + return "cpa:antigravity:credits-hint:" + strings.TrimSpace(authID) +} + func antigravityCreditsAvailableForModel(auth *Auth, model string) bool { if auth == nil { return false diff --git a/sdk/cliproxy/auth/antigravity_credits_test.go b/sdk/cliproxy/auth/antigravity_credits_test.go index 59d5aaa6..540a4ef0 100644 --- a/sdk/cliproxy/auth/antigravity_credits_test.go +++ b/sdk/cliproxy/auth/antigravity_credits_test.go @@ -9,6 +9,7 @@ import ( "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" log "github.com/sirupsen/logrus" @@ -127,6 +128,35 @@ func TestManagerExecuteStream_AntigravityCreditsFallbackAfterBootstrap429(t *tes } } +func TestManagerExecuteStream_AntigravityCreditsHomeKVUnavailableFailsRequest(t *testing.T) { + const model = "claude-opus-4-6-thinking" + executor := &antigravityCreditsFallbackExecutor{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true}, + }) + manager.RegisterExecutor(executor) + registry.GetGlobalRegistry().RegisterClient("ag-credits-home-kv", "antigravity", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("ag-credits-home-kv") }) + homekv.SetCurrent(homekv.New(internalconfig.HomeConfig{Enabled: false})) + t.Cleanup(homekv.ClearCurrent) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: "ag-credits-home-kv", Provider: "antigravity"}); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + _, errExecute := manager.ExecuteStream(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("ExecuteStream() error = nil, want home kv unavailable error") + } + if status := statusCodeFromError(errExecute); status != http.StatusServiceUnavailable { + t.Fatalf("ExecuteStream() status = %d, want %d; err=%v", status, http.StatusServiceUnavailable, errExecute) + } + if !strings.Contains(errExecute.Error(), "home kv store unavailable") { + t.Fatalf("ExecuteStream() error = %v, want home kv store unavailable", errExecute) + } +} + func TestManagerExecuteStream_CodexOnlyDoesNotEnterAntigravityCreditsFallback(t *testing.T) { const model = "gpt-5.5" logger := log.StandardLogger() diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 08b81dad..78d98eff 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1578,7 +1578,9 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye } if lastErr != nil { if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { - if resp, ok := m.tryAntigravityCreditsExecute(ctx, req, opts); ok { + if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { + return cliproxyexecutor.Response{}, errCredits + } else if ok { return resp, nil } } @@ -1644,7 +1646,9 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli } if lastErr != nil { if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { - if result, ok := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); ok { + if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { + return nil, errCredits + } else if ok { return result, nil } } @@ -4218,15 +4222,13 @@ func requestedModelFromMetadata(metadata map[string]any, fallback string) string return fallback } -func (m *Manager) findAllAntigravityCreditsCandidateAuths(routeModel string, opts cliproxyexecutor.Options) []creditsCandidateEntry { +func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, routeModel string, opts cliproxyexecutor.Options) ([]creditsCandidateEntry, error) { if m == nil { - return nil + return nil, nil } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + var candidates []creditsCandidateEntry m.mu.RLock() - defer m.mu.RUnlock() - var known []creditsCandidateEntry - var unknown []creditsCandidateEntry for _, auth := range m.auths { if auth == nil || auth.Disabled || auth.Status == StatusDisabled { continue @@ -4245,32 +4247,37 @@ func (m *Manager) findAllAntigravityCreditsCandidateAuths(routeModel string, opt if !ok { continue } - - hint, okHint := GetAntigravityCreditsHint(auth.ID) - if okHint && hint.Known { - if !hint.Available { - continue - } - known = append(known, creditsCandidateEntry{ - auth: auth.Clone(), - executor: executor, - provider: providerKey, - }) - continue - } - unknown = append(unknown, creditsCandidateEntry{ + candidates = append(candidates, creditsCandidateEntry{ auth: auth.Clone(), executor: executor, provider: providerKey, }) } + m.mu.RUnlock() + + var known []creditsCandidateEntry + var unknown []creditsCandidateEntry + for _, candidate := range candidates { + hint, okHint, errHint := GetAntigravityCreditsHintRequired(ctx, candidate.auth.ID) + if errHint != nil { + return nil, antigravityCreditsKVUnavailableError(errHint) + } + if okHint && hint.Known { + if !hint.Available { + continue + } + known = append(known, candidate) + continue + } + unknown = append(unknown, candidate) + } sort.Slice(known, func(i, j int) bool { return known[i].auth.ID < known[j].auth.ID }) sort.Slice(unknown, func(i, j int) bool { return unknown[i].auth.ID < unknown[j].auth.ID }) - return append(known, unknown...) + return append(known, unknown...), nil } type creditsCandidateEntry struct { @@ -4320,12 +4327,15 @@ func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, provider } } -func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool) { +func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool, error) { routeModel := req.Model - candidates := m.findAllAntigravityCreditsCandidateAuths(routeModel, opts) + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) + if errCandidates != nil { + return cliproxyexecutor.Response{}, false, errCandidates + } for _, c := range candidates { if ctx.Err() != nil { - return cliproxyexecutor.Response{}, false + return cliproxyexecutor.Response{}, false, nil } creditsCtx := WithAntigravityCredits(ctx) if rt := m.roundTripperFor(c.auth); rt != nil { @@ -4362,18 +4372,21 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy continue } m.MarkResult(creditsCtx, result) - return resp, true + return resp, true, nil } } - return cliproxyexecutor.Response{}, false + return cliproxyexecutor.Response{}, false, nil } -func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool) { +func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool, error) { routeModel := req.Model - candidates := m.findAllAntigravityCreditsCandidateAuths(routeModel, opts) + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) + if errCandidates != nil { + return nil, false, errCandidates + } for _, c := range candidates { if ctx.Err() != nil { - return nil, false + return nil, false, nil } creditsCtx := WithAntigravityCredits(ctx) if rt := m.roundTripperFor(c.auth); rt != nil { @@ -4395,9 +4408,16 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl if errStream != nil { continue } - return result, true + return result, true, nil } - return nil, false + return nil, false, nil +} + +func antigravityCreditsKVUnavailableError(cause error) error { + if cause == nil { + return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable: " + cause.Error(), HTTPStatus: http.StatusServiceUnavailable} } func (m *Manager) persist(ctx context.Context, auth *Auth) error { diff --git a/sdk/cliproxy/auth/conductor_credits_candidates_test.go b/sdk/cliproxy/auth/conductor_credits_candidates_test.go index f9487b0b..ade8e6b4 100644 --- a/sdk/cliproxy/auth/conductor_credits_candidates_test.go +++ b/sdk/cliproxy/auth/conductor_credits_candidates_test.go @@ -1,9 +1,14 @@ package auth import ( + "context" + "net/http" + "strings" "testing" "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) @@ -32,7 +37,10 @@ func TestFindAllAntigravityCreditsCandidateAuths_PrefersKnownCreditsThenUnknown( opts := cliproxyexecutor.Options{} - candidates := m.findAllAntigravityCreditsCandidateAuths("claude-sonnet-4-6", opts) + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(context.Background(), "claude-sonnet-4-6", opts) + if errCandidates != nil { + t.Fatalf("findAllAntigravityCreditsCandidateAuths() error = %v", errCandidates) + } if len(candidates) != 2 { t.Fatalf("candidates len = %d, want 2", len(candidates)) } @@ -43,7 +51,10 @@ func TestFindAllAntigravityCreditsCandidateAuths_PrefersKnownCreditsThenUnknown( t.Fatalf("candidates[1].auth.ID = %q, want %q", candidates[1].auth.ID, "aa-unknown") } - nonClaude := m.findAllAntigravityCreditsCandidateAuths("gemini-3-flash", opts) + nonClaude, errNonClaude := m.findAllAntigravityCreditsCandidateAuths(context.Background(), "gemini-3-flash", opts) + if errNonClaude != nil { + t.Fatalf("findAllAntigravityCreditsCandidateAuths(non claude) error = %v", errNonClaude) + } if len(nonClaude) != 0 { t.Fatalf("nonClaude len = %d, want 0", len(nonClaude)) } @@ -51,7 +62,10 @@ func TestFindAllAntigravityCreditsCandidateAuths_PrefersKnownCreditsThenUnknown( pinnedOpts := cliproxyexecutor.Options{ Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: "aa-unknown"}, } - pinned := m.findAllAntigravityCreditsCandidateAuths("claude-sonnet-4-6", pinnedOpts) + pinned, errPinned := m.findAllAntigravityCreditsCandidateAuths(context.Background(), "claude-sonnet-4-6", pinnedOpts) + if errPinned != nil { + t.Fatalf("findAllAntigravityCreditsCandidateAuths(pinned) error = %v", errPinned) + } if len(pinned) != 1 { t.Fatalf("pinned len = %d, want 1", len(pinned)) } @@ -59,3 +73,28 @@ func TestFindAllAntigravityCreditsCandidateAuths_PrefersKnownCreditsThenUnknown( t.Fatalf("pinned[0].auth.ID = %q, want %q", pinned[0].auth.ID, "aa-unknown") } } + +func TestFindAllAntigravityCreditsCandidateAuths_HomeKVUnavailableReturnsError(t *testing.T) { + homekv.SetCurrent(homekv.New(internalconfig.HomeConfig{Enabled: false})) + t.Cleanup(homekv.ClearCurrent) + + m := &Manager{ + auths: map[string]*Auth{ + "ag-home-kv": {ID: "ag-home-kv", Provider: "antigravity"}, + }, + executors: map[string]ProviderExecutor{ + "antigravity": schedulerTestExecutor{}, + }, + } + + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(context.Background(), "claude-sonnet-4-6", cliproxyexecutor.Options{}) + if errCandidates == nil { + t.Fatalf("findAllAntigravityCreditsCandidateAuths() error = nil, candidates=%#v", candidates) + } + if status := statusCodeFromError(errCandidates); status != http.StatusServiceUnavailable { + t.Fatalf("statusCodeFromError() = %d, want %d; err=%v", status, http.StatusServiceUnavailable, errCandidates) + } + if !strings.Contains(errCandidates.Error(), "home kv store unavailable") { + t.Fatalf("error = %v, want home kv store unavailable", errCandidates) + } +}