mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 01:16:47 +08:00
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:
+89
-5
@@ -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
@@ -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)
|
||||
|
||||
Vendored
+108
-11
@@ -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
@@ -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("")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user