From d8f2dceef789feb5dd565b4868021c78e30003aa Mon Sep 17 00:00:00 2001 From: rome-xi <49893941+rome-xi@users.noreply.github.com> Date: Mon, 7 Sep 2026 01:26:52 +0000 Subject: [PATCH] perf(antigravity): batch functionResponse name repairs Repair placeholder functionResponse names with part-local sjson and one Index splice of the request body, matching the sibling content-edit path instead of one full-body sjson.SetBytes per repaired name. --- .../runtime/executor/antigravity_executor.go | 59 ++++-- ...y_preupstream_rewrite_differential_test.go | 194 +++++++++++++++++- 2 files changed, 237 insertions(+), 16 deletions(-) diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 29c129062..9c9435499 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -523,6 +523,7 @@ func ensureAntigravityGeminiBoundaryUserContent(modelName string, payload []byte type antigravityContentEdit struct { index int64 + path string start int end int replacement []byte @@ -651,6 +652,13 @@ func normalizeAntigravityGeminiFunctionResponseRoles(rawJSON []byte) []byte { }) return true }) + return applyAntigravityIndexedEdits(rawJSON, edits, validOffsets) +} + +// applyAntigravityIndexedEdits splices collected JSON fragments into the original +// request with one body copy. Applying SJSON once per field made large histories +// scale with history size multiplied by the number of edits. +func applyAntigravityIndexedEdits(rawJSON []byte, edits []antigravityContentEdit, validOffsets bool) []byte { if len(edits) == 0 { return rawJSON } @@ -685,7 +693,10 @@ func normalizeAntigravityGeminiFunctionResponseRoles(rawJSON []byte) []byte { func applyAntigravityContentEditsWithSJSON(rawJSON []byte, edits []antigravityContentEdit) []byte { out := rawJSON for _, edit := range edits { - path := fmt.Sprintf("request.contents.%d", edit.index) + path := edit.path + if path == "" { + path = fmt.Sprintf("request.contents.%d", edit.index) + } if updated, errSet := sjson.SetRawBytes(out, path, edit.replacement); errSet == nil { out = updated } @@ -693,6 +704,9 @@ func applyAntigravityContentEditsWithSJSON(rawJSON []byte, edits []antigravityCo return out } +// repairAntigravityGeminiFunctionResponseNames copies missing or placeholder +// functionResponse names from the matching functionCall. Edits are applied to +// each part in isolation, then spliced into the request with one body copy. func repairAntigravityGeminiFunctionResponseNames(rawJSON []byte) []byte { contents := util.GetGJSONBytesNoCopy(rawJSON, "request.contents") if !contents.IsArray() { @@ -721,7 +735,8 @@ func repairAntigravityGeminiFunctionResponseNames(rawJSON []byte) []byte { return rawJSON } - out := rawJSON + edits := make([]antigravityContentEdit, 0) + validOffsets := true contents.ForEach(func(contentIdx, content gjson.Result) bool { parts := content.Get("parts") if !parts.IsArray() { @@ -729,23 +744,39 @@ func repairAntigravityGeminiFunctionResponseNames(rawJSON []byte) []byte { } parts.ForEach(func(partIdx, part gjson.Result) bool { fr := part.Get("functionResponse") - if fr.Exists() { - id := strings.TrimSpace(fr.Get("id").String()) - name := strings.TrimSpace(fr.Get("name").String()) - if id != "" && (name == "" || name == "unknown") { - if realName, ok := callIDToName[id]; ok { - path := fmt.Sprintf("request.contents.%d.parts.%d.functionResponse.name", contentIdx.Int(), partIdx.Int()) - if updated, errSet := sjson.SetBytes(out, path, realName); errSet == nil { - out = updated - } - } - } + if !fr.Exists() { + return true } + id := strings.TrimSpace(fr.Get("id").String()) + name := strings.TrimSpace(fr.Get("name").String()) + if id == "" || (name != "" && name != "unknown") { + return true + } + realName, ok := callIDToName[id] + if !ok { + return true + } + updatedPart, errSet := sjson.SetBytes([]byte(part.Raw), "functionResponse.name", realName) + if errSet != nil { + return true + } + start := part.Index + end := start + len(part.Raw) + if start < 0 || end < start || end > len(rawJSON) || !bytes.Equal(rawJSON[start:end], []byte(part.Raw)) { + validOffsets = false + } + edits = append(edits, antigravityContentEdit{ + index: contentIdx.Int(), + path: fmt.Sprintf("request.contents.%d.parts.%d", contentIdx.Int(), partIdx.Int()), + start: start, + end: end, + replacement: updatedPart, + }) return true }) return true }) - return out + return applyAntigravityIndexedEdits(rawJSON, edits, validOffsets) } func validateAntigravityRequestSignatures(ctx context.Context, modelName string, from sdktranslator.Format, rawJSON []byte) ([]byte, error) { diff --git a/internal/runtime/executor/antigravity_preupstream_rewrite_differential_test.go b/internal/runtime/executor/antigravity_preupstream_rewrite_differential_test.go index 1f4da536f..126b45161 100644 --- a/internal/runtime/executor/antigravity_preupstream_rewrite_differential_test.go +++ b/internal/runtime/executor/antigravity_preupstream_rewrite_differential_test.go @@ -5,9 +5,11 @@ import ( "encoding/json" "fmt" "math/rand" + "strings" "testing" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -16,6 +18,8 @@ func TestNormalizeAntigravityGeminiFunctionResponseRolesMatchesLegacy(t *testing []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read","args":{}}},{"functionCall":{"id":"call-2","name":"write","args":{}}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-2","name":"write","response":{"ok":2}}},{"functionResponse":{"id":"call-1","name":"read","response":{"ok":1}}}]}]}}`), []byte("{\r\n \"request\" : {\r\n \"contents\" : [\r\n {\"role\":\"model\",\"parts\":[{\"functionCall\":{\"id\":\"a\",\"name\":\"one\"}},{\"functionCall\":{\"id\":\"b\",\"name\":\"two\"}}]},\r\n {\"role\" : \"user\", \"parts\" : [ { \"functionResponse\" : {\"id\":\"a\",\"name\":\"one\"} }, { \"functionResponse\" : {\"id\":\"b\",\"name\":\"two\"} } ]}\r\n ]\r\n }\r\n}"), []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"actual"}}]},{"parts":[{"functionResponse":{"id":"a","name":"unknown"}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"actual"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","name":"","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"actual"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","response":{"ok":true}}}]}]}}`), []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"one"}}]},{"role":"user","role":"model","parts":[{"functionResponse":{"id":"a","name":"one"}}]}]}}`), []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"one"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","name":"one"}}],"parts":[{"text":"duplicate"}]}]}}`), []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"one"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","name":"one"}}]}],"contents":[{"role":"user","parts":[]}]}}`), @@ -71,8 +75,15 @@ func randomAntigravityFunctionHistory(randomSource *rand.Rand) []byte { call["id"] = id response["id"] = id } - if id != "" && randomSource.Intn(8) == 0 { - response["name"] = "unknown" + if id != "" { + switch randomSource.Intn(10) { + case 0: + response["name"] = "unknown" + case 1: + response["name"] = "" + case 2: + delete(response, "name") + } } calls = append(calls, map[string]any{"functionCall": call}) responses = append(responses, map[string]any{"functionResponse": response}) @@ -104,6 +115,142 @@ func randomAntigravityFunctionHistory(randomSource *rand.Rand) []byte { return payload } +func TestRepairAntigravityGeminiFunctionResponseNamesMatchesLegacy(t *testing.T) { + fixtures := [][]byte{ + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"read","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"name":"unknown","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown","response":{"ok":true}}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}},{"functionCall":{"id":"call-2","name":"write"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown"}},{"functionResponse":{"id":"call-2","name":""}}]}]}}`), + []byte("{\r\n \"request\" : {\r\n \"contents\" : [\r\n {\"role\":\"model\",\"parts\":[{\"functionCall\":{\"id\":\"a\",\"name\":\"one\"}}]},\r\n {\"role\" : \"user\", \"parts\" : [ { \"functionResponse\" : {\"id\":\"a\",\"name\":\"unknown\"} } ]}\r\n ]\r\n }\r\n}"), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"unknown"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","name":"unknown"}}]}]}}`), + []byte(`{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"a","name":"one"}}]},{"role":"user","parts":[{"functionResponse":{"id":"a","name":"one"}}]}]`), + } + + randomSource := rand.New(rand.NewSource(0xA617A5)) + for range 1_000 { + fixtures = append(fixtures, randomAntigravityFunctionHistory(randomSource)) + } + + changed := 0 + unchanged := 0 + for index, fixture := range fixtures { + want := legacyRepairAntigravityGeminiFunctionResponseNames(fixture) + got := repairAntigravityGeminiFunctionResponseNames(fixture) + if !bytes.Equal(got, want) { + t.Fatalf("case %d differs: input_bytes=%d got_bytes=%d want_bytes=%d got=%s want=%s", index, len(fixture), len(got), len(want), got, want) + } + if bytes.Equal(fixture, want) { + unchanged++ + } else { + changed++ + } + if again := repairAntigravityGeminiFunctionResponseNames(got); !bytes.Equal(again, got) { + t.Fatalf("case %d is not idempotent", index) + } + } + if changed == 0 || unchanged == 0 { + t.Fatalf("degenerate fixtures: changed=%d unchanged=%d", changed, unchanged) + } +} + +func TestRepairAntigravityGeminiFunctionResponseNames(t *testing.T) { + tests := []struct { + name string + payload string + wantName map[string]string + wantUnchanged bool + }{ + { + name: "empty name repaired", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "read"}, + }, + { + name: "unknown name repaired", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "read"}, + }, + { + name: "missing name field repaired", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "read"}, + }, + { + name: "good name left alone", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"already_good","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "already_good"}, + wantUnchanged: true, + }, + { + name: "missing id skipped", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}}]},{"role":"user","parts":[{"functionResponse":{"name":"unknown","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "unknown"}, + wantUnchanged: true, + }, + { + name: "multiple repairs in one request", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"read"}},{"functionCall":{"id":"call-2","name":"write"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-2","name":"unknown"}},{"functionResponse":{"id":"call-1","name":""}}]}]}}`, + wantName: map[string]string{ + "request.contents.1.parts.0.functionResponse.name": "write", + "request.contents.1.parts.1.functionResponse.name": "read", + }, + }, + { + name: "no-op when map empty", + payload: `{"request":{"contents":[{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown","response":{"ok":true}}}]}]}}`, + wantName: map[string]string{"request.contents.0.parts.0.functionResponse.name": "unknown"}, + wantUnchanged: true, + }, + { + name: "functionCall name unknown is not mapped", + payload: `{"request":{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"unknown"}}]},{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown"}}]}]}}`, + wantName: map[string]string{"request.contents.1.parts.0.functionResponse.name": "unknown"}, + wantUnchanged: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + payload := []byte(test.payload) + got := repairAntigravityGeminiFunctionResponseNames(payload) + wantLegacy := legacyRepairAntigravityGeminiFunctionResponseNames(payload) + if !bytes.Equal(got, wantLegacy) { + t.Fatalf("differs from legacy oracle: got=%s want=%s", got, wantLegacy) + } + if test.wantUnchanged && !bytes.Equal(got, payload) { + t.Fatalf("payload changed: got=%s", got) + } + for path, wantName := range test.wantName { + if gotName := gjson.GetBytes(got, path).String(); gotName != wantName { + t.Fatalf("%s = %q, want %q; output=%s", path, gotName, wantName, got) + } + } + }) + } +} + +func TestApplyAntigravityIndexedEditsFallsBackWhenOffsetsInvalid(t *testing.T) { + payload := []byte(`{"request":{"contents":[{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"unknown"}}]}]}}`) + replacement := []byte(`{"functionResponse":{"id":"call-1","name":"read"}}`) + edits := []antigravityContentEdit{{ + path: "request.contents.0.parts.0", + start: -1, + end: -1, + replacement: replacement, + }} + want, errSet := sjson.SetRawBytes(payload, "request.contents.0.parts.0", replacement) + if errSet != nil { + t.Fatal(errSet) + } + got := applyAntigravityIndexedEdits(payload, edits, false) + if !bytes.Equal(got, want) { + t.Fatalf("fallback differs: got=%s want=%s", got, want) + } +} + func TestApplyAntigravityContentEditsWithSJSONFallback(t *testing.T) { payload := []byte(`{"request":{"contents":[{"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"read"}}]}]}}`) replacement := []byte(`{"role":"model","parts":[{"functionResponse":{"id":"call-1","name":"read"}}]}`) @@ -240,3 +387,46 @@ func randomAntigravitySchemaRequest(randomSource *rand.Rand) string { } return string(payload) } + +var antigravityNameRepairBenchmarkOutput []byte + +func BenchmarkRepairAntigravityGeminiFunctionResponseNames(b *testing.B) { + payload := syntheticAntigravityNameRepairBenchmarkPayload(1<<20, 32) + b.Run("legacy", func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(len(payload))) + for b.Loop() { + antigravityNameRepairBenchmarkOutput = legacyRepairAntigravityGeminiFunctionResponseNames(payload) + } + }) + b.Run("batched", func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(len(payload))) + for b.Loop() { + antigravityNameRepairBenchmarkOutput = repairAntigravityGeminiFunctionResponseNames(payload) + } + }) +} + +func syntheticAntigravityNameRepairBenchmarkPayload(inlineBytes, turns int) []byte { + var payload strings.Builder + payload.Grow(inlineBytes + turns*320) + payload.WriteString(`{"request":{"contents":[{"role":"user","parts":[{"inlineData":{"mimeType":"application/octet-stream","data":"`) + payload.WriteString(strings.Repeat("a", inlineBytes)) + payload.WriteString(`"}}]}`) + for turn := range turns { + fmt.Fprintf( + &payload, + `,{"role":"model","parts":[{"functionCall":{"id":"call-%d","name":"lookup","args":{"turn":%d}}}]}`, + turn, + turn, + ) + fmt.Fprintf( + &payload, + `,{"role":"user","parts":[{"functionResponse":{"id":"call-%d","name":"unknown","response":{"result":"ok"}}}]}`, + turn, + ) + } + payload.WriteString(`]}}`) + return []byte(payload.String()) +}