refactor(translator): enhance signature handling in Claude and Gemini requests, streamline cache usage and remove unnecessary tests

This commit is contained in:
hkfires
2026-01-21 20:21:49 +08:00
parent d9c6317c84
commit c8884f5e25
5 changed files with 52 additions and 90 deletions

2
go.mod
View File

@@ -21,6 +21,7 @@ require (
golang.org/x/crypto v0.45.0 golang.org/x/crypto v0.45.0
golang.org/x/net v0.47.0 golang.org/x/net v0.47.0
golang.org/x/oauth2 v0.30.0 golang.org/x/oauth2 v0.30.0
golang.org/x/text v0.31.0
gopkg.in/natefinch/lumberjack.v2 v2.2.1 gopkg.in/natefinch/lumberjack.v2 v2.2.1
gopkg.in/yaml.v3 v3.0.1 gopkg.in/yaml.v3 v3.0.1
) )
@@ -70,7 +71,6 @@ require (
golang.org/x/arch v0.8.0 // indirect golang.org/x/arch v0.8.0 // indirect
golang.org/x/sync v0.18.0 // indirect golang.org/x/sync v0.18.0 // indirect
golang.org/x/sys v0.38.0 // indirect golang.org/x/sys v0.38.0 // indirect
golang.org/x/text v0.31.0 // indirect
google.golang.org/protobuf v1.34.1 // indirect google.golang.org/protobuf v1.34.1 // indirect
gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/ini.v1 v1.67.0 // indirect
) )

View File

@@ -98,38 +98,32 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
// Use GetThinkingText to handle wrapped thinking objects // Use GetThinkingText to handle wrapped thinking objects
thinkingText := thinking.GetThinkingText(contentResult) thinkingText := thinking.GetThinkingText(contentResult)
// Always try cached signature first (more reliable than client-provided)
// Client may send stale or invalid signatures from different sessions
signature := "" signature := ""
signatureResult := contentResult.Get("signature") if thinkingText != "" {
hasClientSignature := signatureResult.Exists() && signatureResult.String() != "" if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" {
signature = cachedSig
// Only consider cached signatures when the client provided a signature. // log.Debugf("Using cached signature for thinking block")
// Unsigned thinking blocks must be dropped.
if hasClientSignature {
// Always try cached signature first (more reliable than client-provided)
// Client may send stale or invalid signatures from other requests
if thinkingText != "" {
if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" {
signature = cachedSig
// log.Debugf("Using cached signature for thinking block")
}
} }
}
// Fallback to client signature only if cache miss and client signature is valid // Fallback to client signature only if cache miss and client signature is valid
if signature == "" { if signature == "" {
clientSignature := "" signatureResult := contentResult.Get("signature")
if signatureResult.Exists() && signatureResult.String() != "" { clientSignature := ""
arrayClientSignatures := strings.SplitN(signatureResult.String(), "#", 2) if signatureResult.Exists() && signatureResult.String() != "" {
if len(arrayClientSignatures) == 2 { arrayClientSignatures := strings.SplitN(signatureResult.String(), "#", 2)
if modelName == arrayClientSignatures[0] { if len(arrayClientSignatures) == 2 {
clientSignature = arrayClientSignatures[1] if modelName == arrayClientSignatures[0] {
} clientSignature = arrayClientSignatures[1]
} }
} }
if cache.HasValidSignature(modelName, clientSignature) {
signature = clientSignature
}
// log.Debugf("Using client-provided signature for thinking block")
} }
if cache.HasValidSignature(modelName, clientSignature) {
signature = clientSignature
}
// log.Debugf("Using client-provided signature for thinking block")
} }
// Store for subsequent tool_use in the same message // Store for subsequent tool_use in the same message

View File

@@ -74,6 +74,8 @@ func TestConvertClaudeRequestToAntigravity_RoleMapping(t *testing.T) {
} }
func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) { func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) {
cache.ClearSignatureCache("")
// Valid signature must be at least 50 characters // Valid signature must be at least 50 characters
validSignature := "abc123validSignature1234567890123456789012345678901234567890" validSignature := "abc123validSignature1234567890123456789012345678901234567890"
thinkingText := "Let me think..." thinkingText := "Let me think..."
@@ -115,6 +117,8 @@ func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) {
} }
func TestConvertClaudeRequestToAntigravity_ThinkingBlockWithoutSignature(t *testing.T) { func TestConvertClaudeRequestToAntigravity_ThinkingBlockWithoutSignature(t *testing.T) {
cache.ClearSignatureCache("")
// Unsigned thinking blocks should be removed entirely (not converted to text) // Unsigned thinking blocks should be removed entirely (not converted to text)
inputJSON := []byte(`{ inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking", "model": "claude-sonnet-4-5-thinking",
@@ -236,6 +240,8 @@ func TestConvertClaudeRequestToAntigravity_ToolUse(t *testing.T) {
} }
func TestConvertClaudeRequestToAntigravity_ToolUse_WithSignature(t *testing.T) { func TestConvertClaudeRequestToAntigravity_ToolUse_WithSignature(t *testing.T) {
cache.ClearSignatureCache("")
validSignature := "abc123validSignature1234567890123456789012345678901234567890" validSignature := "abc123validSignature1234567890123456789012345678901234567890"
thinkingText := "Let me think..." thinkingText := "Let me think..."
@@ -277,6 +283,8 @@ func TestConvertClaudeRequestToAntigravity_ToolUse_WithSignature(t *testing.T) {
} }
func TestConvertClaudeRequestToAntigravity_ReorderThinking(t *testing.T) { func TestConvertClaudeRequestToAntigravity_ReorderThinking(t *testing.T) {
cache.ClearSignatureCache("")
// Case: text block followed by thinking block -> should be reordered to thinking first // Case: text block followed by thinking block -> should be reordered to thinking first
validSignature := "abc123validSignature1234567890123456789012345678901234567890" validSignature := "abc123validSignature1234567890123456789012345678901234567890"
thinkingText := "Planning..." thinkingText := "Planning..."
@@ -485,6 +493,8 @@ func TestConvertClaudeRequestToAntigravity_TrailingUnsignedThinking_Removed(t *t
} }
func TestConvertClaudeRequestToAntigravity_TrailingSignedThinking_Kept(t *testing.T) { func TestConvertClaudeRequestToAntigravity_TrailingSignedThinking_Kept(t *testing.T) {
cache.ClearSignatureCache("")
// Last assistant message ends with signed thinking block - should be kept // Last assistant message ends with signed thinking block - should be kept
validSignature := "abc123validSignature1234567890123456789012345678901234567890" validSignature := "abc123validSignature1234567890123456789012345678901234567890"
thinkingText := "Valid thinking..." thinkingText := "Valid thinking..."

View File

@@ -99,44 +99,36 @@ func ConvertGeminiRequestToAntigravity(modelName string, inputRawJSON []byte, _
} }
// Gemini-specific handling for non-Claude models: // Gemini-specific handling for non-Claude models:
// - Remove thinking parts entirely.
// - Add skip_thought_signature_validator to functionCall parts so upstream can bypass signature validation. // - Add skip_thought_signature_validator to functionCall parts so upstream can bypass signature validation.
// - Also mark thinking parts with the same sentinel when present (we keep the parts; we only annotate them).
if !strings.Contains(modelName, "claude") { if !strings.Contains(modelName, "claude") {
const skipSentinel = "skip_thought_signature_validator" const skipSentinel = "skip_thought_signature_validator"
gjson.GetBytes(rawJSON, "request.contents").ForEach(func(contentIdx, content gjson.Result) bool { gjson.GetBytes(rawJSON, "request.contents").ForEach(func(contentIdx, content gjson.Result) bool {
if content.Get("role").String() != "model" { if content.Get("role").String() == "model" {
return true // First pass: collect indices of thinking parts to mark with skip sentinel
} var thinkingIndicesToSkipSignature []int64
partsResult := content.Get("parts") content.Get("parts").ForEach(func(partIdx, part gjson.Result) bool {
if !partsResult.IsArray() { // Collect indices of thinking blocks to mark with skip sentinel
return true if part.Get("thought").Bool() {
} thinkingIndicesToSkipSignature = append(thinkingIndicesToSkipSignature, partIdx.Int())
}
parts := partsResult.Array() // Add skip sentinel to functionCall parts
newParts := make([]interface{}, 0, len(parts)) if part.Get("functionCall").Exists() {
for _, part := range parts { existingSig := part.Get("thoughtSignature").String()
if part.Get("thought").Bool() { if existingSig == "" || len(existingSig) < 50 {
continue rawJSON, _ = sjson.SetBytes(rawJSON, fmt.Sprintf("request.contents.%d.parts.%d.thoughtSignature", contentIdx.Int(), partIdx.Int()), skipSentinel)
}
partRaw := part.Raw
if part.Get("functionCall").Exists() {
existingSig := part.Get("thoughtSignature").String()
if existingSig == "" || len(existingSig) < 50 {
updatedPart, errSet := sjson.Set(partRaw, "thoughtSignature", skipSentinel)
if errSet != nil {
log.WithError(errSet).Debug("failed to set thoughtSignature on functionCall part")
} else {
partRaw = updatedPart
} }
} }
return true
})
// Add skip_thought_signature_validator sentinel to thinking blocks in reverse order to preserve indices
for i := len(thinkingIndicesToSkipSignature) - 1; i >= 0; i-- {
idx := thinkingIndicesToSkipSignature[i]
rawJSON, _ = sjson.SetBytes(rawJSON, fmt.Sprintf("request.contents.%d.parts.%d.thoughtSignature", contentIdx.Int(), idx), skipSentinel)
} }
newParts = append(newParts, gjson.Parse(partRaw).Value())
} }
rawJSON, _ = sjson.SetBytes(rawJSON, fmt.Sprintf("request.contents.%d.parts", contentIdx.Int()), newParts)
return true return true
}) })
} }

View File

@@ -62,40 +62,6 @@ func TestConvertGeminiRequestToAntigravity_AddSkipSentinelToFunctionCall(t *test
} }
} }
func TestConvertGeminiRequestToAntigravity_RemoveThinkingBlocks(t *testing.T) {
// Thinking blocks should be removed entirely for Gemini
validSignature := "abc123validSignature1234567890123456789012345678901234567890"
inputJSON := []byte(fmt.Sprintf(`{
"model": "gemini-3-pro-preview",
"contents": [
{
"role": "model",
"parts": [
{"thought": true, "text": "Thinking...", "thoughtSignature": "%s"},
{"text": "Here is my response"}
]
}
]
}`, validSignature))
output := ConvertGeminiRequestToAntigravity("gemini-3-pro-preview", inputJSON, false)
outputStr := string(output)
// Check that thinking block is removed
parts := gjson.Get(outputStr, "request.contents.0.parts").Array()
if len(parts) != 1 {
t.Fatalf("Expected 1 part (thinking removed), got %d", len(parts))
}
// Only text part should remain
if parts[0].Get("thought").Bool() {
t.Error("Thinking block should be removed for Gemini")
}
if parts[0].Get("text").String() != "Here is my response" {
t.Errorf("Expected text 'Here is my response', got '%s'", parts[0].Get("text").String())
}
}
func TestConvertGeminiRequestToAntigravity_ParallelFunctionCalls(t *testing.T) { func TestConvertGeminiRequestToAntigravity_ParallelFunctionCalls(t *testing.T) {
// Multiple functionCalls should all get skip_thought_signature_validator // Multiple functionCalls should all get skip_thought_signature_validator
inputJSON := []byte(`{ inputJSON := []byte(`{