From c8884f5e2588b186fa0a37afa30e27e8582cdaad Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Wed, 21 Jan 2026 20:21:49 +0800 Subject: [PATCH] refactor(translator): enhance signature handling in Claude and Gemini requests, streamline cache usage and remove unnecessary tests --- go.mod | 2 +- .../claude/antigravity_claude_request.go | 46 ++++++++--------- .../claude/antigravity_claude_request_test.go | 10 ++++ .../gemini/antigravity_gemini_request.go | 50 ++++++++----------- .../gemini/antigravity_gemini_request_test.go | 34 ------------- 5 files changed, 52 insertions(+), 90 deletions(-) diff --git a/go.mod b/go.mod index 963d9c49..863d0413 100644 --- a/go.mod +++ b/go.mod @@ -21,6 +21,7 @@ require ( golang.org/x/crypto v0.45.0 golang.org/x/net v0.47.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/yaml.v3 v3.0.1 ) @@ -70,7 +71,6 @@ require ( golang.org/x/arch v0.8.0 // indirect golang.org/x/sync v0.18.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 gopkg.in/ini.v1 v1.67.0 // indirect ) diff --git a/internal/translator/antigravity/claude/antigravity_claude_request.go b/internal/translator/antigravity/claude/antigravity_claude_request.go index bce76892..e87a7d6b 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_request.go +++ b/internal/translator/antigravity/claude/antigravity_claude_request.go @@ -98,38 +98,32 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _ // Use GetThinkingText to handle wrapped thinking objects 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 := "" - signatureResult := contentResult.Get("signature") - hasClientSignature := signatureResult.Exists() && signatureResult.String() != "" - - // Only consider cached signatures when the client provided a signature. - // 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") - } + 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 - if signature == "" { - clientSignature := "" - if signatureResult.Exists() && signatureResult.String() != "" { - arrayClientSignatures := strings.SplitN(signatureResult.String(), "#", 2) - if len(arrayClientSignatures) == 2 { - if modelName == arrayClientSignatures[0] { - clientSignature = arrayClientSignatures[1] - } + // Fallback to client signature only if cache miss and client signature is valid + if signature == "" { + signatureResult := contentResult.Get("signature") + clientSignature := "" + if signatureResult.Exists() && signatureResult.String() != "" { + arrayClientSignatures := strings.SplitN(signatureResult.String(), "#", 2) + if len(arrayClientSignatures) == 2 { + 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 diff --git a/internal/translator/antigravity/claude/antigravity_claude_request_test.go b/internal/translator/antigravity/claude/antigravity_claude_request_test.go index 7831b8bd..9f40b9fa 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_request_test.go +++ b/internal/translator/antigravity/claude/antigravity_claude_request_test.go @@ -74,6 +74,8 @@ func TestConvertClaudeRequestToAntigravity_RoleMapping(t *testing.T) { } func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) { + cache.ClearSignatureCache("") + // Valid signature must be at least 50 characters validSignature := "abc123validSignature1234567890123456789012345678901234567890" thinkingText := "Let me think..." @@ -115,6 +117,8 @@ func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) { } func TestConvertClaudeRequestToAntigravity_ThinkingBlockWithoutSignature(t *testing.T) { + cache.ClearSignatureCache("") + // Unsigned thinking blocks should be removed entirely (not converted to text) inputJSON := []byte(`{ "model": "claude-sonnet-4-5-thinking", @@ -236,6 +240,8 @@ func TestConvertClaudeRequestToAntigravity_ToolUse(t *testing.T) { } func TestConvertClaudeRequestToAntigravity_ToolUse_WithSignature(t *testing.T) { + cache.ClearSignatureCache("") + validSignature := "abc123validSignature1234567890123456789012345678901234567890" thinkingText := "Let me think..." @@ -277,6 +283,8 @@ func TestConvertClaudeRequestToAntigravity_ToolUse_WithSignature(t *testing.T) { } func TestConvertClaudeRequestToAntigravity_ReorderThinking(t *testing.T) { + cache.ClearSignatureCache("") + // Case: text block followed by thinking block -> should be reordered to thinking first validSignature := "abc123validSignature1234567890123456789012345678901234567890" thinkingText := "Planning..." @@ -485,6 +493,8 @@ func TestConvertClaudeRequestToAntigravity_TrailingUnsignedThinking_Removed(t *t } func TestConvertClaudeRequestToAntigravity_TrailingSignedThinking_Kept(t *testing.T) { + cache.ClearSignatureCache("") + // Last assistant message ends with signed thinking block - should be kept validSignature := "abc123validSignature1234567890123456789012345678901234567890" thinkingText := "Valid thinking..." diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_request.go b/internal/translator/antigravity/gemini/antigravity_gemini_request.go index 37346119..2ad9bd80 100644 --- a/internal/translator/antigravity/gemini/antigravity_gemini_request.go +++ b/internal/translator/antigravity/gemini/antigravity_gemini_request.go @@ -99,44 +99,36 @@ func ConvertGeminiRequestToAntigravity(modelName string, inputRawJSON []byte, _ } // 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. + // - Also mark thinking parts with the same sentinel when present (we keep the parts; we only annotate them). if !strings.Contains(modelName, "claude") { const skipSentinel = "skip_thought_signature_validator" gjson.GetBytes(rawJSON, "request.contents").ForEach(func(contentIdx, content gjson.Result) bool { - if content.Get("role").String() != "model" { - return true - } - partsResult := content.Get("parts") - if !partsResult.IsArray() { - return true - } - - parts := partsResult.Array() - newParts := make([]interface{}, 0, len(parts)) - for _, part := range parts { - if part.Get("thought").Bool() { - continue - } - - 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 + if content.Get("role").String() == "model" { + // First pass: collect indices of thinking parts to mark with skip sentinel + var thinkingIndicesToSkipSignature []int64 + content.Get("parts").ForEach(func(partIdx, part gjson.Result) bool { + // Collect indices of thinking blocks to mark with skip sentinel + if part.Get("thought").Bool() { + thinkingIndicesToSkipSignature = append(thinkingIndicesToSkipSignature, partIdx.Int()) + } + // Add skip sentinel to functionCall parts + if part.Get("functionCall").Exists() { + existingSig := part.Get("thoughtSignature").String() + if existingSig == "" || len(existingSig) < 50 { + rawJSON, _ = sjson.SetBytes(rawJSON, fmt.Sprintf("request.contents.%d.parts.%d.thoughtSignature", contentIdx.Int(), partIdx.Int()), skipSentinel) } } + 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 }) } diff --git a/internal/translator/antigravity/gemini/antigravity_gemini_request_test.go b/internal/translator/antigravity/gemini/antigravity_gemini_request_test.go index 58cffd69..8867a30e 100644 --- a/internal/translator/antigravity/gemini/antigravity_gemini_request_test.go +++ b/internal/translator/antigravity/gemini/antigravity_gemini_request_test.go @@ -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) { // Multiple functionCalls should all get skip_thought_signature_validator inputJSON := []byte(`{