diff --git a/internal/translator/antigravity/claude/antigravity_claude_response.go b/internal/translator/antigravity/claude/antigravity_claude_response.go index 52fc358e..d26a1c9f 100644 --- a/internal/translator/antigravity/claude/antigravity_claude_response.go +++ b/internal/translator/antigravity/claude/antigravity_claude_response.go @@ -9,11 +9,16 @@ package claude import ( "bytes" "context" + "crypto/sha256" + "encoding/hex" "fmt" "strings" "sync/atomic" "time" + "github.com/router-for-me/CLIProxyAPI/v6/internal/cache" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -35,6 +40,31 @@ type Params struct { HasSentFinalEvents bool // Indicates if final content/message events have been sent HasToolUse bool // Indicates if tool use was observed in the stream HasContent bool // Tracks whether any content (text, thinking, or tool use) has been output + + // P3: Signature caching support + SessionID string // Session ID derived from request for signature caching + CurrentThinkingText strings.Builder // Accumulates thinking text for signature caching +} + +// deriveSessionIDFromRequest generates a stable session ID from the request JSON. +func deriveSessionIDFromRequest(rawJSON []byte) string { + messages := gjson.GetBytes(rawJSON, "messages") + if !messages.IsArray() { + return "" + } + for _, msg := range messages.Array() { + if msg.Get("role").String() == "user" { + content := msg.Get("content").String() + if content == "" { + content = msg.Get("content.0.text").String() + } + if content != "" { + h := sha256.Sum256([]byte(content)) + return hex.EncodeToString(h[:16]) + } + } + } + return "" } // toolUseIDCounter provides a process-wide unique counter for tool use identifiers. @@ -62,6 +92,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq HasFirstResponse: false, ResponseType: 0, ResponseIndex: 0, + SessionID: deriveSessionIDFromRequest(originalRequestRawJSON), } } @@ -119,11 +150,20 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq // Process thinking content (internal reasoning) if partResult.Get("thought").Bool() { if thoughtSignature := partResult.Get("thoughtSignature"); thoughtSignature.Exists() && thoughtSignature.String() != "" { + log.Debug("Branch: signature_delta") + + if params.SessionID != "" && params.CurrentThinkingText.Len() > 0 { + cache.CacheSignature(params.SessionID, params.CurrentThinkingText.String(), thoughtSignature.String()) + log.Debugf("Cached signature for thinking block (sessionID=%s, textLen=%d)", params.SessionID, params.CurrentThinkingText.Len()) + params.CurrentThinkingText.Reset() + } + 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", 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 + params.CurrentThinkingText.WriteString(partTextResult.String()) output = output + "event: content_block_delta\n" data, _ := sjson.Set(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, params.ResponseIndex), "delta.thinking", partTextResult.String()) output = output + fmt.Sprintf("data: %s\n\n\n", data) @@ -152,6 +192,9 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq output = output + fmt.Sprintf("data: %s\n\n\n", data) params.ResponseType = 2 // Set state to thinking params.HasContent = true + // P3: Start accumulating thinking text for signature caching + params.CurrentThinkingText.Reset() + params.CurrentThinkingText.WriteString(partTextResult.String()) } } else { finishReasonResult := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason") diff --git a/internal/translator/antigravity/claude/antigravity_claude_response_test.go b/internal/translator/antigravity/claude/antigravity_claude_response_test.go new file mode 100644 index 00000000..7ffd7666 --- /dev/null +++ b/internal/translator/antigravity/claude/antigravity_claude_response_test.go @@ -0,0 +1,389 @@ +package claude + +import ( + "context" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/cache" +) + +func TestConvertBashCommandToCmdField(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "basic command to cmd conversion", + input: `{"command": "git diff"}`, + expected: `{"cmd":"git diff"}`, + }, + { + name: "already has cmd field - no change", + input: `{"cmd": "git diff"}`, + expected: `{"cmd": "git diff"}`, + }, + { + name: "both cmd and command - keep cmd only", + input: `{"command": "git diff", "cmd": "ls"}`, + expected: `{"command": "git diff", "cmd": "ls"}`, // no change when cmd exists + }, + { + name: "command with special characters in value", + input: `{"command": "echo \"command\": test"}`, + expected: `{"cmd":"echo \"command\": test"}`, + }, + { + name: "command with nested quotes", + input: `{"command": "bash -c 'echo \"hello\"'"}`, + expected: `{"cmd":"bash -c 'echo \"hello\"'"}`, + }, + { + name: "command with newlines", + input: `{"command": "echo hello\necho world"}`, + expected: `{"cmd":"echo hello\necho world"}`, + }, + { + name: "empty command value", + input: `{"command": ""}`, + expected: `{"cmd":""}`, + }, + { + name: "command with other fields - preserves them", + input: `{"command": "git diff", "timeout": 30}`, + expected: `{ "timeout": 30,"cmd":"git diff"}`, + }, + { + name: "invalid JSON - returns unchanged", + input: `{invalid json`, + expected: `{invalid json`, + }, + { + name: "empty object", + input: `{}`, + expected: `{}`, + }, + { + name: "no command field", + input: `{"restart": true}`, + expected: `{"restart": true}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := convertBashCommandToCmdField(tt.input) + if result != tt.expected { + t.Errorf("convertBashCommandToCmdField(%q) = %q, want %q", tt.input, result, tt.expected) + } + }) + } +} + +// ============================================================================ +// P3: Signature Caching Tests +// ============================================================================ + +func TestConvertAntigravityResponseToClaude_SessionIDDerived(t *testing.T) { + cache.ClearSignatureCache("") + + // Request with user message - should derive session ID + requestJSON := []byte(`{ + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Hello world"}]} + ] + }`) + + // First response chunk with thinking + responseJSON := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "Let me think...", "thought": true}] + } + }] + } + }`) + + var param any + ctx := context.Background() + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, responseJSON, ¶m) + + // Verify session ID was set + params := param.(*Params) + if params.SessionID == "" { + t.Error("SessionID should be derived from request") + } +} + +func TestConvertAntigravityResponseToClaude_ThinkingTextAccumulated(t *testing.T) { + cache.ClearSignatureCache("") + + requestJSON := []byte(`{ + "messages": [{"role": "user", "content": [{"type": "text", "text": "Test"}]}] + }`) + + // First thinking chunk + chunk1 := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "First part of thinking...", "thought": true}] + } + }] + } + }`) + + // Second thinking chunk (continuation) + chunk2 := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": " Second part of thinking...", "thought": true}] + } + }] + } + }`) + + var param any + ctx := context.Background() + + // Process first chunk - starts new thinking block + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk1, ¶m) + params := param.(*Params) + + if params.CurrentThinkingText.Len() == 0 { + t.Error("Thinking text should be accumulated after first chunk") + } + + // Process second chunk - continues thinking block + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk2, ¶m) + + text := params.CurrentThinkingText.String() + if !strings.Contains(text, "First part") || !strings.Contains(text, "Second part") { + t.Errorf("Thinking text should accumulate both parts, got: %s", text) + } +} + +func TestConvertAntigravityResponseToClaude_SignatureCached(t *testing.T) { + cache.ClearSignatureCache("") + + requestJSON := []byte(`{ + "messages": [{"role": "user", "content": [{"type": "text", "text": "Cache test"}]}] + }`) + + // Thinking chunk + thinkingChunk := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "My thinking process here", "thought": true}] + } + }] + } + }`) + + // Signature chunk + validSignature := "abc123validSignature1234567890123456789012345678901234567890" + signatureChunk := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "", "thought": true, "thoughtSignature": "` + validSignature + `"}] + } + }] + } + }`) + + var param any + ctx := context.Background() + + // Process thinking chunk + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, thinkingChunk, ¶m) + params := param.(*Params) + sessionID := params.SessionID + thinkingText := params.CurrentThinkingText.String() + + if sessionID == "" { + t.Fatal("SessionID should be set") + } + if thinkingText == "" { + t.Fatal("Thinking text should be accumulated") + } + + // Process signature chunk - should cache the signature + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, signatureChunk, ¶m) + + // Verify signature was cached + cachedSig := cache.GetCachedSignature(sessionID, thinkingText) + if cachedSig != validSignature { + t.Errorf("Expected cached signature '%s', got '%s'", validSignature, cachedSig) + } + + // Verify thinking text was reset after caching + if params.CurrentThinkingText.Len() != 0 { + t.Error("Thinking text should be reset after signature is cached") + } +} + +func TestConvertAntigravityResponseToClaude_MultipleThinkingBlocks(t *testing.T) { + cache.ClearSignatureCache("") + + requestJSON := []byte(`{ + "messages": [{"role": "user", "content": [{"type": "text", "text": "Multi block test"}]}] + }`) + + validSig1 := "signature1_12345678901234567890123456789012345678901234567" + validSig2 := "signature2_12345678901234567890123456789012345678901234567" + + // First thinking block with signature + block1Thinking := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "First thinking block", "thought": true}] + } + }] + } + }`) + block1Sig := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "", "thought": true, "thoughtSignature": "` + validSig1 + `"}] + } + }] + } + }`) + + // Text content (breaks thinking) + textBlock := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "Regular text output"}] + } + }] + } + }`) + + // Second thinking block with signature + block2Thinking := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "Second thinking block", "thought": true}] + } + }] + } + }`) + block2Sig := []byte(`{ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "", "thought": true, "thoughtSignature": "` + validSig2 + `"}] + } + }] + } + }`) + + var param any + ctx := context.Background() + + // Process first thinking block + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, block1Thinking, ¶m) + params := param.(*Params) + sessionID := params.SessionID + firstThinkingText := params.CurrentThinkingText.String() + + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, block1Sig, ¶m) + + // Verify first signature cached + if cache.GetCachedSignature(sessionID, firstThinkingText) != validSig1 { + t.Error("First thinking block signature should be cached") + } + + // Process text (transitions out of thinking) + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, textBlock, ¶m) + + // Process second thinking block + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, block2Thinking, ¶m) + secondThinkingText := params.CurrentThinkingText.String() + + ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, block2Sig, ¶m) + + // Verify second signature cached + if cache.GetCachedSignature(sessionID, secondThinkingText) != validSig2 { + t.Error("Second thinking block signature should be cached") + } +} + +func TestDeriveSessionIDFromRequest(t *testing.T) { + tests := []struct { + name string + input []byte + wantEmpty bool + }{ + { + name: "valid user message", + input: []byte(`{"messages": [{"role": "user", "content": "Hello"}]}`), + wantEmpty: false, + }, + { + name: "user message with content array", + input: []byte(`{"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}]}`), + wantEmpty: false, + }, + { + name: "no user message", + input: []byte(`{"messages": [{"role": "assistant", "content": "Hi"}]}`), + wantEmpty: true, + }, + { + name: "empty messages", + input: []byte(`{"messages": []}`), + wantEmpty: true, + }, + { + name: "no messages field", + input: []byte(`{}`), + wantEmpty: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := deriveSessionIDFromRequest(tt.input) + if tt.wantEmpty && result != "" { + t.Errorf("Expected empty session ID, got '%s'", result) + } + if !tt.wantEmpty && result == "" { + t.Error("Expected non-empty session ID") + } + }) + } +} + +func TestDeriveSessionIDFromRequest_Deterministic(t *testing.T) { + input := []byte(`{"messages": [{"role": "user", "content": "Same message"}]}`) + + id1 := deriveSessionIDFromRequest(input) + id2 := deriveSessionIDFromRequest(input) + + if id1 != id2 { + t.Errorf("Session ID should be deterministic: '%s' != '%s'", id1, id2) + } +} + +func TestDeriveSessionIDFromRequest_DifferentMessages(t *testing.T) { + input1 := []byte(`{"messages": [{"role": "user", "content": "Message A"}]}`) + input2 := []byte(`{"messages": [{"role": "user", "content": "Message B"}]}`) + + id1 := deriveSessionIDFromRequest(input1) + id2 := deriveSessionIDFromRequest(input2) + + if id1 == id2 { + t.Error("Different messages should produce different session IDs") + } +}