From 8d9f4edf9b2f6c47824829612e0f2c72014c3be6 Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Tue, 20 Jan 2026 13:45:25 +0800
Subject: [PATCH] feat(translator): unify model group references by introducing
 `GetModelGroup` helper function

---
 internal/cache/signature_cache.go                | 16 ++++++++++++++--
 .../claude/antigravity_claude_response.go        |  4 ++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/internal/cache/signature_cache.go b/internal/cache/signature_cache.go
index ee8ad0b2..477b1245 100644
--- a/internal/cache/signature_cache.go
+++ b/internal/cache/signature_cache.go
@@ -4,6 +4,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 )
@@ -103,7 +104,7 @@ func CacheSignature(modelName, sessionID, text, signature string) {
 		return
 	}
 
-	sc := getOrCreateSession(fmt.Sprintf("%s#%s", modelName, sessionID))
+	sc := getOrCreateSession(fmt.Sprintf("%s#%s", GetModelGroup(modelName), sessionID))
 	textHash := hashText(text)
 
 	sc.mu.Lock()
@@ -122,7 +123,7 @@ func GetCachedSignature(modelName, sessionID, text string) string {
 		return ""
 	}
 
-	val, ok := signatureCache.Load(fmt.Sprintf("%s#%s", modelName, sessionID))
+	val, ok := signatureCache.Load(fmt.Sprintf("%s#%s", GetModelGroup(modelName), sessionID))
 	if !ok {
 		return ""
 	}
@@ -168,3 +169,14 @@ func ClearSignatureCache(sessionID string) {
 func HasValidSignature(signature string) bool {
 	return signature != "" && len(signature) >= MinValidSignatureLen
 }
+
+func GetModelGroup(modelName string) string {
+	if strings.Contains(modelName, "gpt") {
+		return "gpt"
+	} else if strings.Contains(modelName, "claude") {
+		return "claude"
+	} else if strings.Contains(modelName, "gemini") {
+		return "gemini"
+	}
+	return modelName
+}
diff --git a/internal/translator/antigravity/claude/antigravity_claude_response.go b/internal/translator/antigravity/claude/antigravity_claude_response.go
index c32918d6..e360f850 100644
--- a/internal/translator/antigravity/claude/antigravity_claude_response.go
+++ b/internal/translator/antigravity/claude/antigravity_claude_response.go
@@ -146,7 +146,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
 						}
 
 						output = output + "event: content_block_delta\n"
-						data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":""}}`, params.ResponseIndex), "delta.signature", fmt.Sprintf("%s#%s", modelName, thoughtSignature.String()))
+						data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":""}}`, params.ResponseIndex), "delta.signature", fmt.Sprintf("%s#%s", cache.GetModelGroup(modelName), thoughtSignature.String()))
 						output = output + fmt.Sprintf("data: %s\n\n\n", data)
 						params.HasContent = true
 					} else if params.ResponseType == 2 { // Continue existing thinking block if already in thinking state
@@ -438,7 +438,7 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
 		block := `{"type":"thinking","thinking":""}`
 		block, _ = sjson.Set(block, "thinking", thinkingBuilder.String())
 		if thinkingSignature != "" {
-			block, _ = sjson.Set(block, "signature", fmt.Sprintf("%s#%s", modelName, thinkingSignature))
+			block, _ = sjson.Set(block, "signature", fmt.Sprintf("%s#%s", cache.GetModelGroup(modelName), thinkingSignature))
 		}
 		responseJSON, _ = sjson.SetRaw(responseJSON, "content.-1", block)
 		thinkingBuilder.Reset()