diff --git a/internal/runtime/executor/antigravity_reasoning_replay.go b/internal/runtime/executor/antigravity_reasoning_replay.go index 3b21c99e..68bd45d6 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay.go +++ b/internal/runtime/executor/antigravity_reasoning_replay.go @@ -1236,16 +1236,19 @@ func degradeAntigravityClaudeToolProvenanceIDs(payload []byte) ([]byte, int) { if !contents.IsArray() { return payload, 0 } - out := payload - degraded := 0 - for ci, content := range contents.Array() { + type partReplacement struct { + start int + end int + data []byte + } + replacements := make([]partReplacement, 0) + for _, content := range contents.Array() { parts := content.Get("parts") if !parts.IsArray() { continue } seenFunctionCallInTurn := false - for pi, part := range parts.Array() { - partPath := fmt.Sprintf("request.contents.%d.parts.%d", ci, pi) + for _, part := range parts.Array() { if fc := part.Get("functionCall"); fc.Exists() { isFirstFC := !seenFunctionCallInTurn seenFunctionCallInTurn = true @@ -1253,15 +1256,19 @@ func degradeAntigravityClaudeToolProvenanceIDs(payload []byte) ([]byte, int) { if !util.IsGeminiClaudeToolUseID(id) { continue } - out, _ = sjson.SetBytes(out, partPath+".functionCall.id", antigravitySyntheticToolCallID(id)) + updatedPart, _ := sjson.SetBytes([]byte(part.Raw), "functionCall.id", antigravitySyntheticToolCallID(id)) if part.Get("thoughtSignature").Exists() && part.Get("thoughtSignature").String() != "" { if isFirstFC { - out, _ = sjson.SetBytes(out, partPath+".thoughtSignature", internalsignature.GeminiSkipThoughtSignatureValidator) + updatedPart, _ = sjson.SetBytes(updatedPart, "thoughtSignature", internalsignature.GeminiSkipThoughtSignatureValidator) } else { - out, _ = sjson.DeleteBytes(out, partPath+".thoughtSignature") + updatedPart, _ = sjson.DeleteBytes(updatedPart, "thoughtSignature") } } - degraded++ + replacements = append(replacements, partReplacement{ + start: part.Index, + end: part.Index + len(part.Raw), + data: updatedPart, + }) continue } if fr := part.Get("functionResponse"); fr.Exists() { @@ -1269,12 +1276,36 @@ func degradeAntigravityClaudeToolProvenanceIDs(payload []byte) ([]byte, int) { if !util.IsGeminiClaudeToolUseID(id) { continue } - out, _ = sjson.SetBytes(out, partPath+".functionResponse.id", antigravitySyntheticToolCallID(id)) - degraded++ + updatedPart, _ := sjson.SetBytes([]byte(part.Raw), "functionResponse.id", antigravitySyntheticToolCallID(id)) + replacements = append(replacements, partReplacement{ + start: part.Index, + end: part.Index + len(part.Raw), + data: updatedPart, + }) } } } - return out, degraded + if len(replacements) == 0 { + return payload, 0 + } + + // Mutate each small part independently, then copy the full request only once. + outputSize := len(payload) + for _, replacement := range replacements { + outputSize += len(replacement.data) - (replacement.end - replacement.start) + } + out := make([]byte, 0, outputSize) + last := 0 + for _, replacement := range replacements { + if replacement.start < last || replacement.end > len(payload) { + return payload, 0 + } + out = append(out, payload[last:replacement.start]...) + out = append(out, replacement.data...) + last = replacement.end + } + out = append(out, payload[last:]...) + return out, len(replacements) } // antigravityRepairUnsignedFirstFunctionCalls restores Gemini's bypass sentinel on diff --git a/internal/runtime/executor/antigravity_reasoning_replay_test.go b/internal/runtime/executor/antigravity_reasoning_replay_test.go index 94d1fc52..18db2f91 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay_test.go +++ b/internal/runtime/executor/antigravity_reasoning_replay_test.go @@ -1820,6 +1820,38 @@ func TestDegradeAntigravityClaudeToolProvenanceIDs_AllCallsSigned_OnlyFirstGetsB } } +var antigravityDegradeBenchmarkOutput []byte + +func BenchmarkDegradeAntigravityClaudeToolProvenanceIDs(b *testing.B) { + const calls = 217 + var payload strings.Builder + payload.Grow(1<<20 + calls*512) + payload.WriteString(`{"request":{"contents":[{"role":"user","parts":[{"text":"`) + payload.WriteString(strings.Repeat("x", 1<<20)) + payload.WriteString(`"}]}`) + for i := range calls { + id := util.GeminiClaudeToolUseID(fmt.Sprintf("native-%d", i), "Read", `{"file_path":"/tmp/a"}`) + fmt.Fprintf(&payload, `,{"role":"model","parts":[{"thoughtSignature":"signature-%d","functionCall":{"id":"%s","name":"Read","args":{"file_path":"/tmp/a"}}}]}`, i, id) + fmt.Fprintf(&payload, `,{"role":"user","parts":[{"functionResponse":{"id":"%s","name":"Read","response":{"result":"ok"}}}]}`, id) + } + payload.WriteString(`]}}`) + input := []byte(payload.String()) + if _, degraded := degradeAntigravityClaudeToolProvenanceIDs(input); degraded != calls*2 { + b.Fatalf("degraded = %d, want %d", degraded, calls*2) + } + + b.SetBytes(int64(len(input))) + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + var degraded int + antigravityDegradeBenchmarkOutput, degraded = degradeAntigravityClaudeToolProvenanceIDs(input) + if degraded != calls*2 { + b.Fatalf("degraded = %d, want %d", degraded, calls*2) + } + } +} + func TestPrepareAntigravityGeminiReasoningReplay_ReordersPermutedParallelToolResponsesWithDegradedIDs(t *testing.T) { internalcache.ClearAntigravityReasoningReplayCache() t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache)