fix(antigravity): insert replayed functionCall before matching functionResponse

When HOME replay restores a cached function_call_part and the request
already contains the matching functionResponse, insert a synthetic
role=model content with functionCall (and thoughtSignature) immediately
before that response content instead of writing into the same index.

Add regression tests for user/model functionResponse shapes.
This commit is contained in:
sususu98
2026-06-18 19:50:46 +08:00
parent 292456a884
commit b17d29ad35
2 changed files with 125 additions and 12 deletions
@@ -232,10 +232,42 @@ func antigravityAnyKeyExists(existing map[string]bool, keys []string) bool {
}
func antigravityRequestHasMatchingFunctionResponse(payload []byte, itemResult gjson.Result) bool {
callID := strings.TrimSpace(itemResult.Get("call_id").String())
if callID == "" {
_, ok := antigravityFunctionResponseContentIndex(payload, strings.TrimSpace(itemResult.Get("call_id").String()))
if itemResult.Get("call_id").String() == "" {
return true
}
return ok
}
func antigravityFunctionResponseContentIndex(payload []byte, callID string) (int, bool) {
callID = strings.TrimSpace(callID)
if callID == "" {
return -1, false
}
contents := gjson.GetBytes(payload, "request.contents")
if !contents.IsArray() {
return -1, false
}
for i, content := range contents.Array() {
parts := content.Get("parts")
if !parts.IsArray() {
continue
}
for _, part := range parts.Array() {
fr := part.Get("functionResponse")
if fr.Exists() && strings.TrimSpace(fr.Get("id").String()) == callID {
return i, true
}
}
}
return -1, false
}
func antigravityPayloadHasFunctionCallID(payload []byte, callID string) bool {
callID = strings.TrimSpace(callID)
if callID == "" {
return false
}
contents := gjson.GetBytes(payload, "request.contents")
if !contents.IsArray() {
return false
@@ -246,8 +278,8 @@ func antigravityRequestHasMatchingFunctionResponse(payload []byte, itemResult gj
continue
}
for _, part := range parts.Array() {
fr := part.Get("functionResponse")
if fr.Exists() && strings.TrimSpace(fr.Get("id").String()) == callID {
fc := part.Get("functionCall")
if fc.Exists() && strings.TrimSpace(fc.Get("id").String()) == callID {
return true
}
}
@@ -255,6 +287,52 @@ func antigravityRequestHasMatchingFunctionResponse(payload []byte, itemResult gj
return false
}
func insertAntigravityModelFunctionCallBeforeContent(payload []byte, beforeIndex int, name, callID, thoughtSig string, args gjson.Result) ([]byte, bool) {
contents := gjson.GetBytes(payload, "request.contents")
if !contents.IsArray() {
return payload, false
}
arr := contents.Array()
if beforeIndex < 0 || beforeIndex > len(arr) {
return payload, false
}
fc := map[string]any{"name": name}
if callID != "" {
fc["id"] = callID
}
if args.Exists() {
if args.Type == gjson.String {
fc["args"] = args.String()
} else {
var parsed any
if json.Unmarshal([]byte(args.Raw), &parsed) == nil {
fc["args"] = parsed
}
}
}
part := map[string]any{"functionCall": fc}
if thoughtSig != "" {
part["thoughtSignature"] = thoughtSig
}
newContent := map[string]any{
"role": "model",
"parts": []any{part},
}
newArr := make([]any, 0, len(arr)+1)
for i := 0; i < beforeIndex; i++ {
newArr = append(newArr, arr[i].Value())
}
newArr = append(newArr, newContent)
for i := beforeIndex; i < len(arr); i++ {
newArr = append(newArr, arr[i].Value())
}
updated, err := sjson.SetBytes(payload, "request.contents", newArr)
if err != nil {
return payload, false
}
return updated, true
}
func antigravityRequestHasThoughtSignatureAt(payload []byte, itemResult gjson.Result) bool {
ci := int(itemResult.Get("contentIndex").Int())
pi := int(itemResult.Get("partIndex").Int())
@@ -297,9 +375,22 @@ func insertAntigravityReasoningReplayItems(payload []byte, items [][]byte) ([]by
}
func mergeAntigravityFunctionCallPartReplay(payload []byte, itemResult gjson.Result) ([]byte, bool) {
name := strings.TrimSpace(itemResult.Get("name").String())
args := itemResult.Get("args")
callID := strings.TrimSpace(itemResult.Get("call_id").String())
sig := strings.TrimSpace(itemResult.Get("thoughtSignature").String())
if name == "" || !args.Exists() {
return payload, false
}
if callID != "" && antigravityPayloadHasFunctionCallID(payload, callID) {
return payload, false
}
if frIndex, ok := antigravityFunctionResponseContentIndex(payload, callID); callID != "" && ok {
return insertAntigravityModelFunctionCallBeforeContent(payload, frIndex, name, callID, sig, args)
}
ci := antigravityReasoningReplayResolveContentIndex(payload, int(itemResult.Get("contentIndex").Int()))
pi := int(itemResult.Get("partIndex").Int())
sig := strings.TrimSpace(itemResult.Get("thoughtSignature").String())
pathSig := fmt.Sprintf("request.contents.%d.parts.%d.thoughtSignature", ci, pi)
out := payload
changed := false
@@ -309,11 +400,8 @@ func mergeAntigravityFunctionCallPartReplay(payload []byte, itemResult gjson.Res
changed = true
}
}
name := strings.TrimSpace(itemResult.Get("name").String())
args := itemResult.Get("args")
callID := strings.TrimSpace(itemResult.Get("call_id").String())
pathFC := fmt.Sprintf("request.contents.%d.parts.%d.functionCall", ci, pi)
if !gjson.GetBytes(out, pathFC).Exists() && name != "" && args.Exists() {
if !gjson.GetBytes(out, pathFC).Exists() {
fc := map[string]any{"name": name}
if callID != "" {
fc["id"] = callID
@@ -62,11 +62,36 @@ func TestPrepareAntigravityGeminiReasoningReplayPayloadInjectsCachedToolPart(t *
if !scope.valid() {
t.Fatal("scope invalid")
}
path := "request.contents.1.parts.0.thoughtSignature"
if got := gjson.GetBytes(out, path).String(); got != "sig-first" {
t.Fatalf("%s = %q, want sig-first; body=%s", path, got, string(out))
if gjson.GetBytes(out, "request.contents.1.role").String() != "model" {
t.Fatalf("functionCall replay must be model role at [1], got %s", string(out))
}
if got := gjson.GetBytes(out, "request.contents.1.parts.0.thoughtSignature").String(); got != "sig-first" {
t.Fatalf("thoughtSignature = %q, want sig-first", got)
}
if !gjson.GetBytes(out, "request.contents.1.parts.0.functionCall").Exists() {
t.Fatalf("functionCall not injected: %s", string(out))
}
if !gjson.GetBytes(out, "request.contents.2.parts.0.functionResponse").Exists() {
t.Fatalf("functionResponse should follow model functionCall at [2]: %s", string(out))
}
}
func TestPrepareAntigravityGeminiReasoningReplayInsertsBeforeModelFunctionResponse(t *testing.T) {
internalcache.ClearAntigravityReasoningReplayCache()
t.Cleanup(internalcache.ClearAntigravityReasoningReplayCache)
item := []byte(`{"type":"function_call_part","contentIndex":1,"partIndex":0,"name":"Read","call_id":"id1","args":{"file_path":"/a"},"thoughtSignature":"sig-first"}`)
internalcache.CacheAntigravityReasoningReplayItems("gemini-3-flash-agent", "session:sess-3", [][]byte{item})
payload := []byte(`{"sessionId":"sess-3","request":{"contents":[{"role":"user","parts":[{"text":"hi"}]},{"role":"model","parts":[{"functionResponse":{"id":"id1","name":"Read","response":{"result":"ok"}}}]}]}}`)
out, _, err := prepareAntigravityGeminiReasoningReplayPayload(context.Background(), "gemini-3-flash-agent", cliproxyexecutor.Request{}, cliproxyexecutor.Options{}, payload)
if err != nil {
t.Fatal(err)
}
if !gjson.GetBytes(out, "request.contents.1.parts.0.functionCall").Exists() || gjson.GetBytes(out, "request.contents.1.role").String() != "model" {
t.Fatalf("want model functionCall at [1]: %s", string(out))
}
if !gjson.GetBytes(out, "request.contents.2.parts.0.functionResponse").Exists() {
t.Fatalf("functionResponse should be at [2]: %s", string(out))
}
}