From 40f4df07b7c81a84fe676f3cb15c8fd7c11b50ec Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Mon, 5 Oct 2026 13:24:20 +0800 Subject: [PATCH] feat(antigravity): enhance context hash rejection logging and batch retry logic --- .../executor/antigravity_reasoning_replay.go | 94 +++++++++---- ...antigravity_reasoning_replay_index_test.go | 126 +++++++++++++++++- 2 files changed, 196 insertions(+), 24 deletions(-) diff --git a/internal/runtime/executor/antigravity_reasoning_replay.go b/internal/runtime/executor/antigravity_reasoning_replay.go index 15a6067b1..964284e79 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay.go +++ b/internal/runtime/executor/antigravity_reasoning_replay.go @@ -356,7 +356,13 @@ func applyAntigravityReasoningReplayItems(payload []byte, items [][]byte, toolSc tracker.setPhase(antigravityReplayPhaseBatch) batch := newAntigravityReplayBatch(index) handled := 0 - for handled < len(items) && batch.apply(items[handled], toolSchemas) { + contextHashStopped := false + for handled < len(items) { + rejectionsBefore := tracker.contextHashRejectionCount() + if !batch.apply(items[handled], toolSchemas) { + contextHashStopped = tracker.contextHashRejectionCount() > rejectionsBefore + break + } handled++ } if handled > 0 { @@ -376,10 +382,13 @@ func applyAntigravityReasoningReplayItems(payload []byte, items [][]byte, toolSc break } index = newAntigravityReplayRequestIndexWithTracker(updated, tracker) - // Retry the first unhandled item against the freshly flushed payload. - // It may only have rejected the batch because a prior stable-ID restore - // made the old context index stale. - continue + // Retry only when this batch restored a stable ID. That changes the + // context fingerprint, so the rejection may belong to the stale index. + // A context-hash mismatch on an unchanged identity cannot be repaired + // by rebuilding the same payload. + if !(contextHashStopped && !batch.identityChanged) { + continue + } } // Apply one structural or context-dependent item with the original @@ -1332,17 +1341,30 @@ type antigravityReplayContextHashSample struct { opaqueID bool } -type antigravityReplayContextHashLogTracker struct { - phase antigravityReplayPhase +type antigravityReplayContextHashGroupKey struct { + toolName string + opaqueID bool +} + +type antigravityReplayContextHashGroup struct { + key antigravityReplayContextHashGroupKey batchRejections int sequentialRejections int loggedKeys map[antigravityReplayPartKey]bool - sample *antigravityReplayContextHashSample + sampleContentIndex int + samplePartIndex int +} + +type antigravityReplayContextHashLogTracker struct { + phase antigravityReplayPhase + totalContextHashRejections int + groups map[antigravityReplayContextHashGroupKey]*antigravityReplayContextHashGroup + order []antigravityReplayContextHashGroupKey } func newAntigravityReplayContextHashLogTracker() *antigravityReplayContextHashLogTracker { return &antigravityReplayContextHashLogTracker{ - loggedKeys: make(map[antigravityReplayPartKey]bool), + groups: make(map[antigravityReplayContextHashGroupKey]*antigravityReplayContextHashGroup), } } @@ -1352,33 +1374,59 @@ func (t *antigravityReplayContextHashLogTracker) setPhase(phase antigravityRepla } } +func (t *antigravityReplayContextHashLogTracker) contextHashRejectionCount() int { + if t == nil { + return 0 + } + return t.totalContextHashRejections +} + func (t *antigravityReplayContextHashLogTracker) recordRejection(sample antigravityReplayContextHashSample, key antigravityReplayPartKey) { if t == nil { return } + t.totalContextHashRejections++ + groupKey := antigravityReplayContextHashGroupKey{ + toolName: sample.name, + opaqueID: sample.opaqueID, + } + group, exists := t.groups[groupKey] + if !exists { + group = &antigravityReplayContextHashGroup{ + key: groupKey, + loggedKeys: make(map[antigravityReplayPartKey]bool), + sampleContentIndex: sample.contentIndex, + samplePartIndex: sample.partIndex, + } + t.groups[groupKey] = group + t.order = append(t.order, groupKey) + } if t.phase == antigravityReplayPhaseSequential { - t.sequentialRejections++ + group.sequentialRejections++ } else { - t.batchRejections++ + group.batchRejections++ } - if t.sample == nil { - t.sample = &sample - } - t.loggedKeys[key] = true + group.loggedKeys[key] = true } func (t *antigravityReplayContextHashLogTracker) log() { - if t == nil || t.sample == nil || (t.batchRejections == 0 && t.sequentialRejections == 0) { + if t == nil || len(t.groups) == 0 { return } - partsCount := len(t.loggedKeys) - partsNoun := "parts" - if partsCount == 1 { - partsNoun = "part" + for _, groupKey := range t.order { + group := t.groups[groupKey] + if group == nil || (group.batchRejections == 0 && group.sequentialRejections == 0) { + continue + } + partsCount := len(group.loggedKeys) + partsNoun := "parts" + if partsCount == 1 { + partsNoun = "part" + } + log.Debugf("antigravity replay: context hash rejected %d %s (first=%q at contents[%d].parts[%d], opaque_id=%t, batch=%d, seq=%d)", + partsCount, partsNoun, group.key.toolName, group.sampleContentIndex, group.samplePartIndex, group.key.opaqueID, + group.batchRejections, group.sequentialRejections) } - log.Debugf("antigravity replay: context hash rejected %d %s (first=%q at contents[%d].parts[%d], opaque_id=%t, batch=%d, seq=%d)", - partsCount, partsNoun, t.sample.name, t.sample.contentIndex, t.sample.partIndex, t.sample.opaqueID, - t.batchRejections, t.sequentialRejections) } type antigravityReplayRequestIndex struct { diff --git a/internal/runtime/executor/antigravity_reasoning_replay_index_test.go b/internal/runtime/executor/antigravity_reasoning_replay_index_test.go index 4fc80a81d..f3d33ddbe 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay_index_test.go +++ b/internal/runtime/executor/antigravity_reasoning_replay_index_test.go @@ -1064,6 +1064,53 @@ func TestApplyAntigravityReasoningReplayItems_MultiPartContextHashRejectionLoggi t.Fatal("expected payload not to change") } + execCount := 0 + readCount := 0 + for _, entry := range hook.AllEntries() { + if entry.Level != log.DebugLevel { + continue + } + if strings.Contains(entry.Message, `first="exec"`) { + execCount++ + expected := `antigravity replay: context hash rejected 1 part (first="exec" at contents[1].parts[0], opaque_id=false, batch=1, seq=1)` + if entry.Message != expected { + t.Fatalf("expected exec log %q, got %q", expected, entry.Message) + } + } + if strings.Contains(entry.Message, `first="read"`) { + readCount++ + expected := `antigravity replay: context hash rejected 1 part (first="read" at contents[3].parts[0], opaque_id=false, batch=1, seq=1)` + if entry.Message != expected { + t.Fatalf("expected read log %q, got %q", expected, entry.Message) + } + } + } + + if execCount != 1 || readCount != 1 { + t.Fatalf("expected 1 exec log and 1 read log for different tools, got exec=%d read=%d", execCount, readCount) + } +} + +func TestApplyAntigravityReasoningReplayItems_SameToolMultiPartContextHashRejectionAggregates(t *testing.T) { + hook := newSignatureDebugHook(t) + + payload := []byte(`{"request":{"contents":[ + {"role":"user","parts":[{"text":"first prompt"}]}, + {"role":"model","parts":[{"functionCall":{"id":"call-1","name":"exec","args":{"cmd":"ls"}}}]}, + {"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"exec","response":{"result":"ok"}}}]}, + {"role":"model","parts":[{"functionCall":{"id":"call-2","name":"exec","args":{"cmd":"pwd"}}}]}, + {"role":"user","parts":[{"functionResponse":{"id":"call-2","name":"exec","response":{"result":"/app"}}}]} + ]}}`) + + item1 := []byte(`{"type":"function_call_part","call_id":"call-1","name":"exec","args":{"cmd":"ls"},"thoughtSignature":"sig-1","contextHash":"wrong-hash-1"}`) + item2 := []byte(`{"type":"function_call_part","call_id":"call-2","name":"exec","args":{"cmd":"pwd"},"thoughtSignature":"sig-2","contextHash":"wrong-hash-2"}`) + + toolSchemas := map[string]any{"exec": map[string]any{"type": "object"}} + _, changed := applyAntigravityReasoningReplayItems(payload, [][]byte{item1, item2}, toolSchemas) + if changed { + t.Fatal("expected payload not to change") + } + detailCount := 0 for _, entry := range hook.AllEntries() { if entry.Level != log.DebugLevel { @@ -1079,7 +1126,49 @@ func TestApplyAntigravityReasoningReplayItems_MultiPartContextHashRejectionLoggi } if detailCount != 1 { - t.Fatalf("expected exactly 1 multi-part rejection log, got %d", detailCount) + t.Fatalf("expected exactly 1 aggregated log for same tool across multiple parts, got %d", detailCount) + } +} + +func TestApplyAntigravityReasoningReplayItems_SameToolDifferentOpaqueIDTypesRemainSeparate(t *testing.T) { + hook := newSignatureDebugHook(t) + + nativeID := "call-2" + name := "exec" + args := `{"cmd":"pwd"}` + opaqueID := util.GeminiClaudeToolUseID(nativeID, name, args) + payload := []byte(fmt.Sprintf(`{"request":{"contents":[ + {"role":"user","parts":[{"text":"first prompt"}]}, + {"role":"model","parts":[{"functionCall":{"id":"call-1","name":%q,"args":{"cmd":"ls"}}}]}, + {"role":"user","parts":[{"functionResponse":{"id":"call-1","name":%q,"response":{"result":"ok"}}}]}, + {"role":"model","parts":[{"functionCall":{"id":%q,"name":%q,"args":%s}}]}, + {"role":"user","parts":[{"functionResponse":{"id":%q,"name":%q,"response":{"result":"/app"}}}]} + ]}}`, name, name, opaqueID, name, args, opaqueID, name)) + + item1 := []byte(`{"type":"function_call_part","call_id":"call-1","name":"exec","args":{"cmd":"ls"},"thoughtSignature":"sig-1","contextHash":"wrong-hash-1"}`) + item2 := []byte(fmt.Sprintf(`{"type":"function_call_part","call_id":%q,"name":%q,"args":%s,"thoughtSignature":"sig-2","contextHash":"wrong-hash-2"}`, nativeID, name, args)) + + _, _ = applyAntigravityReasoningReplayItems(payload, [][]byte{item1, item2}, map[string]any{ + "exec": map[string]any{"type": "object"}, + }) + + entries := make([]string, 0, 2) + for _, entry := range hook.AllEntries() { + if entry.Level == log.DebugLevel && strings.Contains(entry.Message, "context hash rejected") { + entries = append(entries, entry.Message) + } + } + expected := []string{ + `antigravity replay: context hash rejected 1 part (first="exec" at contents[1].parts[0], opaque_id=false, batch=1, seq=1)`, + `antigravity replay: context hash rejected 1 part (first="exec" at contents[3].parts[0], opaque_id=true, batch=1, seq=0)`, + } + if len(entries) != len(expected) { + t.Fatalf("expected %d separate logs, got %d: %v", len(expected), len(entries), entries) + } + for index, got := range entries { + if got != expected[index] { + t.Fatalf("log[%d] = %q, want %q", index, got, expected[index]) + } } } @@ -1118,3 +1207,38 @@ func TestApplyAntigravityReasoningReplayItems_SingleContextHashRejectionLogs(t * t.Fatalf("expected exactly 1 detailed rejection log, got %d", detailCount) } } + +func TestApplyAntigravityReasoningReplayItems_ContextHashRejectionSkipsBatchRetry(t *testing.T) { + hook := newSignatureDebugHook(t) + + payload := []byte(`{"request":{"contents":[ + {"role":"user","parts":[{"text":"first prompt"}]}, + {"role":"model","parts":[{"functionCall":{"id":"call-ok","name":"exec","args":{"cmd":"pwd"}},"thoughtSignature":"sig-ok"}]}, + {"role":"user","parts":[{"functionResponse":{"id":"call-ok","name":"exec","response":{"result":"/app"}}}]}, + {"role":"model","parts":[{"functionCall":{"id":"call-1","name":"tool_call","args":{"id":"web_search"}}}]}, + {"role":"user","parts":[{"functionResponse":{"id":"call-1","name":"tool_call","response":{"result":"ok"}}}]} + ]}}`) + + applied := []byte(`{"type":"function_call_part","call_id":"call-ok","name":"exec","args":{"cmd":"pwd"},"thoughtSignature":"sig-ok"}`) + rejected := []byte(`{"type":"function_call_part","call_id":"call-1","name":"tool_call","args":{"id":"web_search"},"thoughtSignature":"sig-1","contextHash":"wrong-hash-1"}`) + toolSchemas := map[string]any{ + "exec": map[string]any{"type": "object"}, + "tool_call": map[string]any{"type": "object"}, + } + _, _ = applyAntigravityReasoningReplayItems(payload, [][]byte{applied, rejected}, toolSchemas) + + detailCount := 0 + for _, entry := range hook.AllEntries() { + if entry.Level != log.DebugLevel || !strings.Contains(entry.Message, "context hash rejected") { + continue + } + detailCount++ + expected := `antigravity replay: context hash rejected 1 part (first="tool_call" at contents[3].parts[0], opaque_id=false, batch=1, seq=1)` + if entry.Message != expected { + t.Fatalf("expected log %q, got %q", expected, entry.Message) + } + } + if detailCount != 1 { + t.Fatalf("expected exactly 1 rejection log, got %d", detailCount) + } +}