feat(antigravity): enhance context hash rejection logging and batch retry logic

2026-10-17 21:19:36 +08:00 · 2026-10-05 13:24:20 +08:00
parent 17dc1b81c2
commit 40f4df07b7
2 changed files with 196 additions and 24 deletions
--- 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 {
--- 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)
+	}
+}