refactor(cache): simplify signature caching by removing sessionID parameter

This commit is contained in:
Luis Pater
2026-01-21 12:38:05 +08:00
parent 88bf4e77ec
commit 65ad5c0c9d
4 changed files with 44 additions and 55 deletions

View File

@@ -96,15 +96,15 @@ func purgeExpiredSessions() {
// CacheSignature stores a thinking signature for a given session and text. // CacheSignature stores a thinking signature for a given session and text.
// Used for Claude models that require signed thinking blocks in multi-turn conversations. // Used for Claude models that require signed thinking blocks in multi-turn conversations.
func CacheSignature(modelName, sessionID, text, signature string) { func CacheSignature(modelName, text, signature string) {
if sessionID == "" || text == "" || signature == "" { if text == "" || signature == "" {
return return
} }
if len(signature) < MinValidSignatureLen { if len(signature) < MinValidSignatureLen {
return return
} }
sc := getOrCreateSession(fmt.Sprintf("%s#%s", GetModelGroup(modelName), sessionID)) sc := getOrCreateSession(fmt.Sprintf("%s#%s", GetModelGroup(modelName), text))
textHash := hashText(text) textHash := hashText(text)
sc.mu.Lock() sc.mu.Lock()
@@ -118,17 +118,17 @@ func CacheSignature(modelName, sessionID, text, signature string) {
// GetCachedSignature retrieves a cached signature for a given session and text. // GetCachedSignature retrieves a cached signature for a given session and text.
// Returns empty string if not found or expired. // Returns empty string if not found or expired.
func GetCachedSignature(modelName, sessionID, text string) string { func GetCachedSignature(modelName, text string) string {
family := GetModelGroup(modelName) family := GetModelGroup(modelName)
if sessionID == "" || text == "" { if text == "" {
if family == "gemini" { if family == "gemini" {
return "skip_thought_signature_validator" return "skip_thought_signature_validator"
} }
return "" return ""
} }
val, ok := signatureCache.Load(fmt.Sprintf("%s#%s", family, sessionID)) val, ok := signatureCache.Load(fmt.Sprintf("%s#%s", family, text))
if !ok { if !ok {
if family == "gemini" { if family == "gemini" {
return "skip_thought_signature_validator" return "skip_thought_signature_validator"

View File

@@ -8,15 +8,14 @@ import (
func TestCacheSignature_BasicStorageAndRetrieval(t *testing.T) { func TestCacheSignature_BasicStorageAndRetrieval(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sessionID := "test-session-1"
text := "This is some thinking text content" text := "This is some thinking text content"
signature := "abc123validSignature1234567890123456789012345678901234567890" signature := "abc123validSignature1234567890123456789012345678901234567890"
// Store signature // Store signature
CacheSignature("test-model", sessionID, text, signature) CacheSignature("test-model", text, signature)
// Retrieve signature // Retrieve signature
retrieved := GetCachedSignature("test-model", sessionID, text) retrieved := GetCachedSignature("test-model", text)
if retrieved != signature { if retrieved != signature {
t.Errorf("Expected signature '%s', got '%s'", signature, retrieved) t.Errorf("Expected signature '%s', got '%s'", signature, retrieved)
} }
@@ -29,13 +28,13 @@ func TestCacheSignature_DifferentSessions(t *testing.T) {
sig1 := "signature1_1234567890123456789012345678901234567890123456" sig1 := "signature1_1234567890123456789012345678901234567890123456"
sig2 := "signature2_1234567890123456789012345678901234567890123456" sig2 := "signature2_1234567890123456789012345678901234567890123456"
CacheSignature("test-model", "session-a", text, sig1) CacheSignature("test-model", text, sig1)
CacheSignature("test-model", "session-b", text, sig2) CacheSignature("test-model", text, sig2)
if GetCachedSignature("test-model", "session-a", text) != sig1 { if GetCachedSignature("test-model", text) != sig1 {
t.Error("Session-a signature mismatch") t.Error("Session-a signature mismatch")
} }
if GetCachedSignature("test-model", "session-b", text) != sig2 { if GetCachedSignature("test-model", text) != sig2 {
t.Error("Session-b signature mismatch") t.Error("Session-b signature mismatch")
} }
} }
@@ -44,13 +43,13 @@ func TestCacheSignature_NotFound(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
// Non-existent session // Non-existent session
if got := GetCachedSignature("test-model", "nonexistent", "some text"); got != "" { if got := GetCachedSignature("test-model", "some text"); got != "" {
t.Errorf("Expected empty string for nonexistent session, got '%s'", got) t.Errorf("Expected empty string for nonexistent session, got '%s'", got)
} }
// Existing session but different text // Existing session but different text
CacheSignature("test-model", "session-x", "text-a", "sigA12345678901234567890123456789012345678901234567890") CacheSignature("test-model", "text-a", "sigA12345678901234567890123456789012345678901234567890")
if got := GetCachedSignature("test-model", "session-x", "text-b"); got != "" { if got := GetCachedSignature("test-model", "text-b"); got != "" {
t.Errorf("Expected empty string for different text, got '%s'", got) t.Errorf("Expected empty string for different text, got '%s'", got)
} }
} }
@@ -59,12 +58,12 @@ func TestCacheSignature_EmptyInputs(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
// All empty/invalid inputs should be no-ops // All empty/invalid inputs should be no-ops
CacheSignature("test-model", "", "text", "sig12345678901234567890123456789012345678901234567890") CacheSignature("test-model", "text", "sig12345678901234567890123456789012345678901234567890")
CacheSignature("test-model", "session", "", "sig12345678901234567890123456789012345678901234567890") CacheSignature("test-model", "", "sig12345678901234567890123456789012345678901234567890")
CacheSignature("test-model", "session", "text", "") CacheSignature("test-model", "text", "")
CacheSignature("test-model", "session", "text", "short") // Too short CacheSignature("test-model", "text", "short") // Too short
if got := GetCachedSignature("test-model", "session", "text"); got != "" { if got := GetCachedSignature("test-model", "text"); got != "" {
t.Errorf("Expected empty after invalid cache attempts, got '%s'", got) t.Errorf("Expected empty after invalid cache attempts, got '%s'", got)
} }
} }
@@ -72,13 +71,12 @@ func TestCacheSignature_EmptyInputs(t *testing.T) {
func TestCacheSignature_ShortSignatureRejected(t *testing.T) { func TestCacheSignature_ShortSignatureRejected(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sessionID := "test-short-sig"
text := "Some text" text := "Some text"
shortSig := "abc123" // Less than 50 chars shortSig := "abc123" // Less than 50 chars
CacheSignature("test-model", sessionID, text, shortSig) CacheSignature("test-model", text, shortSig)
if got := GetCachedSignature("test-model", sessionID, text); got != "" { if got := GetCachedSignature("test-model", text); got != "" {
t.Errorf("Short signature should be rejected, got '%s'", got) t.Errorf("Short signature should be rejected, got '%s'", got)
} }
} }
@@ -87,15 +85,15 @@ func TestClearSignatureCache_SpecificSession(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sig := "validSig1234567890123456789012345678901234567890123456" sig := "validSig1234567890123456789012345678901234567890123456"
CacheSignature("test-model", "session-1", "text", sig) CacheSignature("test-model", "text", sig)
CacheSignature("test-model", "session-2", "text", sig) CacheSignature("test-model", "text", sig)
ClearSignatureCache("session-1") ClearSignatureCache("session-1")
if got := GetCachedSignature("test-model", "session-1", "text"); got != "" { if got := GetCachedSignature("test-model", "text"); got != "" {
t.Error("session-1 should be cleared") t.Error("session-1 should be cleared")
} }
if got := GetCachedSignature("test-model", "session-2", "text"); got != sig { if got := GetCachedSignature("test-model", "text"); got != sig {
t.Error("session-2 should still exist") t.Error("session-2 should still exist")
} }
} }
@@ -104,15 +102,15 @@ func TestClearSignatureCache_AllSessions(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sig := "validSig1234567890123456789012345678901234567890123456" sig := "validSig1234567890123456789012345678901234567890123456"
CacheSignature("test-model", "session-1", "text", sig) CacheSignature("test-model", "text", sig)
CacheSignature("test-model", "session-2", "text", sig) CacheSignature("test-model", "text", sig)
ClearSignatureCache("") ClearSignatureCache("")
if got := GetCachedSignature("test-model", "session-1", "text"); got != "" { if got := GetCachedSignature("test-model", "text"); got != "" {
t.Error("session-1 should be cleared") t.Error("session-1 should be cleared")
} }
if got := GetCachedSignature("test-model", "session-2", "text"); got != "" { if got := GetCachedSignature("test-model", "text"); got != "" {
t.Error("session-2 should be cleared") t.Error("session-2 should be cleared")
} }
} }
@@ -143,21 +141,19 @@ func TestHasValidSignature(t *testing.T) {
func TestCacheSignature_TextHashCollisionResistance(t *testing.T) { func TestCacheSignature_TextHashCollisionResistance(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sessionID := "hash-test-session"
// Different texts should produce different hashes // Different texts should produce different hashes
text1 := "First thinking text" text1 := "First thinking text"
text2 := "Second thinking text" text2 := "Second thinking text"
sig1 := "signature1_1234567890123456789012345678901234567890123456" sig1 := "signature1_1234567890123456789012345678901234567890123456"
sig2 := "signature2_1234567890123456789012345678901234567890123456" sig2 := "signature2_1234567890123456789012345678901234567890123456"
CacheSignature("test-model", sessionID, text1, sig1) CacheSignature("test-model", text1, sig1)
CacheSignature("test-model", sessionID, text2, sig2) CacheSignature("test-model", text2, sig2)
if GetCachedSignature("test-model", sessionID, text1) != sig1 { if GetCachedSignature("test-model", text1) != sig1 {
t.Error("text1 signature mismatch") t.Error("text1 signature mismatch")
} }
if GetCachedSignature("test-model", sessionID, text2) != sig2 { if GetCachedSignature("test-model", text2) != sig2 {
t.Error("text2 signature mismatch") t.Error("text2 signature mismatch")
} }
} }
@@ -165,13 +161,12 @@ func TestCacheSignature_TextHashCollisionResistance(t *testing.T) {
func TestCacheSignature_UnicodeText(t *testing.T) { func TestCacheSignature_UnicodeText(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sessionID := "unicode-session"
text := "한글 텍스트와 이모지 🎉 그리고 特殊文字" text := "한글 텍스트와 이모지 🎉 그리고 特殊文字"
sig := "unicodeSig123456789012345678901234567890123456789012345" sig := "unicodeSig123456789012345678901234567890123456789012345"
CacheSignature("test-model", sessionID, text, sig) CacheSignature("test-model", text, sig)
if got := GetCachedSignature("test-model", sessionID, text); got != sig { if got := GetCachedSignature("test-model", text); got != sig {
t.Errorf("Unicode text signature retrieval failed, got '%s'", got) t.Errorf("Unicode text signature retrieval failed, got '%s'", got)
} }
} }
@@ -179,15 +174,14 @@ func TestCacheSignature_UnicodeText(t *testing.T) {
func TestCacheSignature_Overwrite(t *testing.T) { func TestCacheSignature_Overwrite(t *testing.T) {
ClearSignatureCache("") ClearSignatureCache("")
sessionID := "overwrite-session"
text := "Same text" text := "Same text"
sig1 := "firstSignature12345678901234567890123456789012345678901" sig1 := "firstSignature12345678901234567890123456789012345678901"
sig2 := "secondSignature1234567890123456789012345678901234567890" sig2 := "secondSignature1234567890123456789012345678901234567890"
CacheSignature("test-model", sessionID, text, sig1) CacheSignature("test-model", text, sig1)
CacheSignature("test-model", sessionID, text, sig2) // Overwrite CacheSignature("test-model", text, sig2) // Overwrite
if got := GetCachedSignature("test-model", sessionID, text); got != sig2 { if got := GetCachedSignature("test-model", text); got != sig2 {
t.Errorf("Expected overwritten signature '%s', got '%s'", sig2, got) t.Errorf("Expected overwritten signature '%s', got '%s'", sig2, got)
} }
} }
@@ -199,14 +193,13 @@ func TestCacheSignature_ExpirationLogic(t *testing.T) {
// This test verifies the expiration check exists // This test verifies the expiration check exists
// In a real scenario, we'd mock time.Now() // In a real scenario, we'd mock time.Now()
sessionID := "expiration-test"
text := "text" text := "text"
sig := "validSig1234567890123456789012345678901234567890123456" sig := "validSig1234567890123456789012345678901234567890123456"
CacheSignature("test-model", sessionID, text, sig) CacheSignature("test-model", text, sig)
// Fresh entry should be retrievable // Fresh entry should be retrievable
if got := GetCachedSignature("test-model", sessionID, text); got != sig { if got := GetCachedSignature("test-model", text); got != sig {
t.Errorf("Fresh entry should be retrievable, got '%s'", got) t.Errorf("Fresh entry should be retrievable, got '%s'", got)
} }

View File

@@ -36,11 +36,7 @@ func deriveSessionID(rawJSON []byte) string {
} }
for _, msg := range messages.Array() { for _, msg := range messages.Array() {
if msg.Get("role").String() == "user" { if msg.Get("role").String() == "user" {
content := msg.Get("content").String() content := msg.Get("content").Raw
if content == "" {
// Try to get text from content array
content = msg.Get("content.0.text").String()
}
if content != "" { if content != "" {
h := sha256.Sum256([]byte(content)) h := sha256.Sum256([]byte(content))
return hex.EncodeToString(h[:16]) return hex.EncodeToString(h[:16])
@@ -138,7 +134,7 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
// Client may send stale or invalid signatures from different sessions // Client may send stale or invalid signatures from different sessions
signature := "" signature := ""
if sessionID != "" && thinkingText != "" { if sessionID != "" && thinkingText != "" {
if cachedSig := cache.GetCachedSignature(modelName, sessionID, thinkingText); cachedSig != "" { if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" {
signature = cachedSig signature = cachedSig
// log.Debugf("Using cached signature for thinking block") // log.Debugf("Using cached signature for thinking block")
} }

View File

@@ -140,7 +140,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
// log.Debug("Branch: signature_delta") // log.Debug("Branch: signature_delta")
if params.SessionID != "" && params.CurrentThinkingText.Len() > 0 { if params.SessionID != "" && params.CurrentThinkingText.Len() > 0 {
cache.CacheSignature(modelName, params.SessionID, params.CurrentThinkingText.String(), thoughtSignature.String()) cache.CacheSignature(modelName, params.CurrentThinkingText.String(), thoughtSignature.String())
// log.Debugf("Cached signature for thinking block (sessionID=%s, textLen=%d)", params.SessionID, params.CurrentThinkingText.Len()) // log.Debugf("Cached signature for thinking block (sessionID=%s, textLen=%d)", params.SessionID, params.CurrentThinkingText.Len())
params.CurrentThinkingText.Reset() params.CurrentThinkingText.Reset()
} }