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.
This commit is contained in:
Luis Pater
2026-06-14 21:11:35 +08:00
parent f6aa9a067e
commit 2a050dc95d
29 changed files with 2991 additions and 161 deletions
+89 -5
View File
@@ -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 {
+176
View File
@@ -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)
+108 -11
View File
@@ -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
+200
View File
@@ -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("")
+188
View File
@@ -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
+243
View File
@@ -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
}
+189
View File
@@ -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]
}
+110
View File
@@ -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)
}
}
+236 -28
View File
@@ -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 {
@@ -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) {
@@ -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)
}
+56 -20
View File
@@ -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.
@@ -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 {
+57 -25
View File
@@ -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) {
@@ -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 {
@@ -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)
}
@@ -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)
}
}
@@ -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
@@ -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)
}
}
@@ -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)
}
@@ -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)
}
}
@@ -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)
}
@@ -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")
}
}
@@ -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 {
@@ -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)
+28 -4
View File
@@ -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
@@ -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()
+52 -32
View File
@@ -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 {
@@ -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)
}
}