feat(keyviz): address PR #720 round-1 (Gemini HIGH/MEDIUM + Codex P2)

Three review items, two of them resolved by the same
semantic change (per-cell wire format):

- Gemini HIGH (proto/admin.proto:116) — row-level scalar
  raft_group_id / leader_term insufficient for the
  per-cell dedupe goal. The pivot in admin_grpc.go and
  keyviz_handler.go only captures the FIRST column's
  identity, so subsequent columns under a mid-window
  leader flip get the wrong term and the aggregator's
  (bucket, group, term, column) dedupe key breaks.
- Gemini MEDIUM (keyviz_fanout.go:529) — first-seen
  identity vs max-merge value can mismatch when the
  source whose value won did not match the source whose
  identity was seeded first.

Both resolved by switching wire format to per-cell
parallel arrays:

  // proto/admin.proto KeyVizRow:
  repeated uint64 raft_group_ids = 13;
  repeated uint64 leader_terms = 14;

  // internal/admin JSON KeyVizRow:
  RaftGroupIDs []uint64 `json:"raft_group_ids,omitempty"`
  LeaderTerms  []uint64 `json:"leader_terms,omitempty"`

The pivot loops in admin_grpc.go and keyviz_handler.go
now stamp row.RaftGroupIDs[j] = mr.RaftGroupID per
column (parallel to row.Values[j] = pick(mr)). The
mergeRowInto in keyviz_fanout.go stamps the destination
cell's identity from whichever source contributed the
value kept at that cell — for maxMerge that's the
larger source; for sumMerge that's best-effort
last-touched (sum has no single owner, but in the
steady state both sources agree on the term, so the
distinction doesn't matter).

- Codex P2 (main.go:1565) — publishLeaderTerms reads
  rt.engine twice (nil-check + Status() call) which
  races rt.Close() setting rt.engine = nil during
  startup-error teardown. Snapshotted into a local once.

Caller audit (per loop instructions): the only
production consumers of KeyVizRow.RaftGroupID/LeaderTerm
were the pivot/merge paths in adapter/admin_grpc.go and
internal/admin/keyviz_{handler,fanout}.go, all updated
in this commit. No external consumer reads these fields.
The PR-3c aggregator that will USE the per-cell dedupe
hasn't been written yet, so there are no callers
expecting the old scalar shape.

Test updates:

- TestKeyVizHandlerStampsRaftIdentity: extended to a
  two-column scenario with a mid-window flip; asserts
  parallel arrays match (route 1: groups=[7,7],
  terms=[42,43]; route 2 absent in col1: zero pair).
- TestGetKeyVizMatrixStampsRaftIdentity: same shape on
  the gRPC path against pb.KeyVizRow.RaftGroupIds /
  LeaderTerms.
- TestMergeKeyVizMatricesPreservesRaftIdentity: source
  rows now use parallel-array shape; pinned identity
  survives mergeRowInto.
- TestMergeKeyVizMatricesPerCellIdentityMatchesValueOwner
  (NEW): leadership flip across two cells with each
  leader winning one cell; asserts merged identity at
  each cell matches the value's owner (Gemini MEDIUM
  regression guard).

Sampler tests at the package level are unchanged because
keyviz.MatrixRow still carries scalar RaftGroupID /
LeaderTerm — the per-cell representation only emerges
at the pivot where multiple per-column MatrixRows are
folded into one row-major KeyVizRow.

Author: bootjp

adapter/admin_grpc.go

@@ -612,6 +612,12 @@ func matrixToProto(cols []keyviz.MatrixColumn, pick func(keyviz.MatrixRow) uint6
 				order = append(order, mr.RouteID)
 			}
 			pr.Values[j] = pick(mr)
+			// Per-cell Raft identity stamped from this column's
+			// MatrixRow (Gemini HIGH on PR #720 — a row-level
+			// scalar would only capture the first column's
+			// identity and break the per-cell dedupe goal).
+			pr.RaftGroupIds[j] = mr.RaftGroupID
+			pr.LeaderTerms[j] = mr.LeaderTerm
 		}
 	}
 	resp.Rows = make([]*pb.KeyVizRow, len(order))
@@ -675,9 +681,12 @@ func newKeyVizRowFrom(mr keyviz.MatrixRow, numCols int) *pb.KeyVizRow {
 		Aggregate:         mr.Aggregate,
 		RouteCount:        total,
 		RouteIdsTruncated: mr.Aggregate && total > uint64(len(mr.MemberRoutes)),
-		RaftGroupId:       mr.RaftGroupID,
-		LeaderTerm:        mr.LeaderTerm,
 		Values:            make([]uint64, numCols),
+		// Per-cell parallel arrays — populated in the column-loop
+		// above, not here, because the first MatrixRow seen for a
+		// route only carries that one column's identity.
+		RaftGroupIds: make([]uint64, numCols),
+		LeaderTerms:  make([]uint64, numCols),
 	}
 	if mr.Aggregate {
 		row.RouteIds = append([]uint64(nil), mr.MemberRoutes...)

adapter/admin_grpc_keyviz_test.go

@@ -274,14 +274,18 @@ func TestGetKeyVizMatrixHonorsRowsBudget(t *testing.T) {
 	require.Equal(t, "route:4", resp.Rows[1].BucketId)
 }
 
-// TestGetKeyVizMatrixStampsRaftIdentity pins the Phase 2-C+ wire
-// extension: MatrixRow.RaftGroupID and MatrixRow.LeaderTerm propagate
-// through matrixToProto into the proto KeyVizRow's
-// raft_group_id (field 13) and leader_term (field 14). The fan-out
-// aggregator's per-term dedupe key requires both fields on the wire.
+// TestGetKeyVizMatrixStampsRaftIdentity pins the Phase 2-C+ per-cell
+// wire extension: MatrixRow.RaftGroupID and MatrixRow.LeaderTerm
+// propagate per-column through matrixToProto into the proto
+// KeyVizRow's raft_group_ids (field 13) and leader_terms (field 14)
+// parallel-to-values arrays. The fan-out aggregator's per-cell
+// dedupe key requires per-column identity (Gemini HIGH on PR #720
+// — row-level scalars only captured the first column's identity
+// and broke dedupe under mid-window leader flips).
 func TestGetKeyVizMatrixStampsRaftIdentity(t *testing.T) {
 	t.Parallel()
 	t0 := time.Unix(1_700_000_000, 0)
+	t1 := t0.Add(time.Minute)
 	srv := newAdminServerWithFakeSampler(t, []keyviz.MatrixColumn{
 		{
 			At: t0,
@@ -290,19 +294,27 @@ func TestGetKeyVizMatrixStampsRaftIdentity(t *testing.T) {
 				{RouteID: 2, Start: []byte("b"), End: []byte("c"), Writes: 9, RaftGroupID: 0, LeaderTerm: 0},
 			},
 		},
+		// col 1: term flipped on route 1 mid-window; route 2 absent.
+		{
+			At: t1,
+			Rows: []keyviz.MatrixRow{
+				{RouteID: 1, Start: []byte("a"), End: []byte("b"), Writes: 11, RaftGroupID: 7, LeaderTerm: 43},
+			},
+		},
 	})
 
 	resp, err := srv.GetKeyVizMatrix(context.Background(), &pb.GetKeyVizMatrixRequest{
 		Series: pb.KeyVizSeries_KEYVIZ_SERIES_WRITES,
 	})
 	require.NoError(t, err)
 	require.Len(t, resp.Rows, 2)
-	// route:1 — non-zero identity propagated.
+	// route:1 — per-cell identity captures the term flip.
 	require.Equal(t, "route:1", resp.Rows[0].BucketId)
-	require.Equal(t, uint64(7), resp.Rows[0].RaftGroupId, "RaftGroupID must propagate to proto field 13")
-	require.Equal(t, uint64(42), resp.Rows[0].LeaderTerm, "LeaderTerm must propagate to proto field 14")
-	// route:2 — zero values stay zero (legacy max-merge fallback).
+	require.Equal(t, []uint64{7, 7}, resp.Rows[0].RaftGroupIds, "raft_group_ids must be parallel to values")
+	require.Equal(t, []uint64{42, 43}, resp.Rows[0].LeaderTerms,
+		"leader_terms must capture per-column term so the aggregator can dedupe across the flip")
+	// route:2 — present only in col0; col1 zero is the "term not tracked" sentinel.
 	require.Equal(t, "route:2", resp.Rows[1].BucketId)
-	require.Equal(t, uint64(0), resp.Rows[1].RaftGroupId)
-	require.Equal(t, uint64(0), resp.Rows[1].LeaderTerm)
+	require.Equal(t, []uint64{0, 0}, resp.Rows[1].RaftGroupIds)
+	require.Equal(t, []uint64{0, 0}, resp.Rows[1].LeaderTerms)
 }

internal/admin/keyviz_fanout.go

@@ -525,9 +525,14 @@ func mergeRowInto(
 			RouteIDs:          append([]uint64(nil), row.RouteIDs...),
 			RouteIDsTruncated: row.RouteIDsTruncated,
 			RouteCount:        row.RouteCount,
-			RaftGroupID:       row.RaftGroupID,
-			LeaderTerm:        row.LeaderTerm,
 			Values:            make([]uint64, mergedWidth),
+			// Per-cell parallel arrays sized to the merged column
+			// width. Stamped per-cell in the loop below from the
+			// largest-source for that cell, so the identity always
+			// matches the value we kept (Gemini MEDIUM on PR #720
+			// resolved by going per-cell).
+			RaftGroupIDs: make([]uint64, mergedWidth),
+			LeaderTerms:  make([]uint64, mergedWidth),
 		}
 		rowsByBucket[row.BucketID] = dst
 		*bucketOrder = append(*bucketOrder, row.BucketID)
@@ -537,11 +542,26 @@ func mergeRowInto(
 		if !ok || j >= len(row.Values) {
 			continue
 		}
-		next, conflict := mergeFn(dst.Values[idx], row.Values[j])
+		prev := dst.Values[idx]
+		next, conflict := mergeFn(prev, row.Values[j])
 		dst.Values[idx] = next
 		if conflict {
 			dst.Conflict = true
 		}
+		// Identity belongs to the source whose value we kept. For
+		// sumMerge (reads) this is best-effort: prev+incoming has
+		// no single owner, so we keep whichever side contributed
+		// most recently — close-to-correct in the steady state.
+		// For maxMerge (writes), `next == row.Values[j]` exactly
+		// when the incoming source won the cell; in the tied
+		// (prev == incoming != 0) case `next == prev`, both sides
+		// agree on the value, and either identity is acceptable.
+		if next == row.Values[j] && j < len(row.RaftGroupIDs) {
+			dst.RaftGroupIDs[idx] = row.RaftGroupIDs[j]
+		}
+		if next == row.Values[j] && j < len(row.LeaderTerms) {
+			dst.LeaderTerms[idx] = row.LeaderTerms[j]
+		}
 	}
 }
 

internal/admin/keyviz_fanout_test.go

@@ -71,33 +71,67 @@ func TestMergeKeyVizMatricesWritesMaxStableLeader(t *testing.T) {
 }
 
 // TestMergeKeyVizMatricesPreservesRaftIdentity pins the Phase 2-C+
-// wire extension on the fan-out merge path: when mergeRowInto seeds
-// the destination row from the first source, RaftGroupID and
-// LeaderTerm are copied through. PR-3c will use these fields to
-// switch from §4.2's row-level max-merge to §9.1's per-cell
-// (group, term) dedupe; this PR ensures the fields survive the
-// merge so PR-3c has data to act on.
+// per-cell wire extension on the fan-out merge path: mergeRowInto
+// stamps the destination row's per-cell (RaftGroupIDs[idx],
+// LeaderTerms[idx]) from whichever source contributed the value
+// kept at that cell. Both sources reporting the same identity
+// for a writes-max merge is the steady-state shape — merged
+// identity matches regardless of source order. (Gemini HIGH on
+// PR #720 resolved by going per-cell; row-level scalars would
+// only capture the first column's identity and break the per-cell
+// dedupe goal.)
 func TestMergeKeyVizMatricesPreservesRaftIdentity(t *testing.T) {
 	t.Parallel()
 	col := []int64{1_700_000_000_000}
 	a := KeyVizMatrix{
 		ColumnUnixMs: col,
 		Series:       keyVizSeriesWrites,
 		Rows: []KeyVizRow{
-			{BucketID: "route:5", Values: []uint64{30}, RaftGroupID: 7, LeaderTerm: 42},
+			{BucketID: "route:5", Values: []uint64{30}, RaftGroupIDs: []uint64{7}, LeaderTerms: []uint64{42}},
 		},
 	}
 	b := KeyVizMatrix{
 		ColumnUnixMs: col,
 		Series:       keyVizSeriesWrites,
 		Rows: []KeyVizRow{
-			{BucketID: "route:5", Values: []uint64{0}, RaftGroupID: 7, LeaderTerm: 42},
+			{BucketID: "route:5", Values: []uint64{0}, RaftGroupIDs: []uint64{7}, LeaderTerms: []uint64{42}},
 		},
 	}
 	merged := mergeKeyVizMatrices([]KeyVizMatrix{a, b}, keyVizSeriesWrites)
 	require.Len(t, merged.Rows, 1)
-	require.Equal(t, uint64(7), merged.Rows[0].RaftGroupID, "RaftGroupID must survive mergeRowInto")
-	require.Equal(t, uint64(42), merged.Rows[0].LeaderTerm, "LeaderTerm must survive mergeRowInto")
+	require.Equal(t, []uint64{7}, merged.Rows[0].RaftGroupIDs, "RaftGroupIDs must survive mergeRowInto")
+	require.Equal(t, []uint64{42}, merged.Rows[0].LeaderTerms, "LeaderTerms must survive mergeRowInto")
+}
+
+// TestMergeKeyVizMatricesPerCellIdentityMatchesValueOwner pins the
+// Gemini MEDIUM fix: when maxMerge picks a value from one source,
+// the identity at that cell must come from the SAME source — not
+// from whoever happened to be processed first. Drives a leadership
+// flip across two consecutive columns with each leader winning
+// one cell.
+func TestMergeKeyVizMatricesPerCellIdentityMatchesValueOwner(t *testing.T) {
+	t.Parallel()
+	col := []int64{1_700_000_000_000, 1_700_000_001_000}
+	exLeader := KeyVizMatrix{
+		ColumnUnixMs: col,
+		Series:       keyVizSeriesWrites,
+		Rows: []KeyVizRow{
+			{BucketID: "route:9", Values: []uint64{50, 0}, RaftGroupIDs: []uint64{7, 7}, LeaderTerms: []uint64{42, 42}},
+		},
+	}
+	newLeader := KeyVizMatrix{
+		ColumnUnixMs: col,
+		Series:       keyVizSeriesWrites,
+		Rows: []KeyVizRow{
+			{BucketID: "route:9", Values: []uint64{0, 80}, RaftGroupIDs: []uint64{7, 7}, LeaderTerms: []uint64{43, 43}},
+		},
+	}
+	merged := mergeKeyVizMatrices([]KeyVizMatrix{exLeader, newLeader}, keyVizSeriesWrites)
+	require.Len(t, merged.Rows, 1)
+	require.Equal(t, []uint64{50, 80}, merged.Rows[0].Values, "writes max-merge picks the larger per cell")
+	require.Equal(t, []uint64{7, 7}, merged.Rows[0].RaftGroupIDs, "groupID stays 7 across both cells")
+	require.Equal(t, []uint64{42, 43}, merged.Rows[0].LeaderTerms,
+		"col0's identity comes from exLeader (term 42, won 50 vs 0); col1's identity comes from newLeader (term 43, won 80 vs 0)")
 }
 
 // TestMergeKeyVizMatricesWritesMaxLeadershipFlip pins §4.2 under a

internal/admin/keyviz_handler.go

@@ -88,17 +88,20 @@ type KeyVizRow struct {
 	RouteCount        uint64   `json:"route_count"`
 	Values            []uint64 `json:"values"`
 	Conflict          bool     `json:"conflict,omitempty"`
-	// RaftGroupID and LeaderTerm carry the route's Raft identity at
-	// flush time. Phase 2-C+ fan-out uses
-	// (bucket_id, raft_group_id, leader_term, column) as the dedupe
-	// key so writes from a leader and the previous leader can be
-	// summed across terms instead of conservatively max-merged.
-	// Zero values mean "term not tracked" (legacy single-group
-	// deployments, virtual aggregate buckets, or nodes that have
-	// not wired the leader-term publisher) — the aggregator falls
-	// back to the legacy max-merge for those cells.
-	RaftGroupID uint64 `json:"raft_group_id,omitempty"`
-	LeaderTerm  uint64 `json:"leader_term,omitempty"`
+	// RaftGroupIDs[j] and LeaderTerms[j] carry the route's Raft
+	// identity at the time column j was flushed (parallel to
+	// Values[]). Phase 2-C+ fan-out uses
+	// (bucket_id, raft_group_id, leader_term, column) as the
+	// dedupe key. Per-cell representation is required because
+	// leadership can flip within the requested window; row-level
+	// scalars would only capture the first column's identity and
+	// cause incorrect dedupe for later columns (Gemini HIGH on
+	// PR #720). Zero values mean "term not tracked" — the
+	// aggregator falls back to the legacy max-merge for those
+	// cells. The slices are either nil (legacy server, no
+	// per-column identity to share) or len == len(Values).
+	RaftGroupIDs []uint64 `json:"raft_group_ids,omitempty"`
+	LeaderTerms  []uint64 `json:"leader_terms,omitempty"`
 	// total accumulates the sum of Values during pivot so the
 	// rowBudget sort is O(N log N) on a precomputed key rather
 	// than O(N log N × M) recomputing the sum per comparison.
@@ -333,6 +336,13 @@ func pivotKeyVizColumns(cols []keyviz.MatrixColumn, series KeyVizSeries, rowBudg
 			}
 			v := pick(mr)
 			row.Values[j] = v
+			// Per-cell Raft identity stamped from this column's
+			// MatrixRow. Each MatrixRow is emitted by Flush() with
+			// the term snapshot taken at that flush moment, so the
+			// j-th column's identity may differ from j-1 if a leader
+			// term flipped between flushes (Gemini HIGH on PR #720).
+			row.RaftGroupIDs[j] = mr.RaftGroupID
+			row.LeaderTerms[j] = mr.LeaderTerm
 			row.total += v
 		}
 	}
@@ -383,9 +393,12 @@ func newKeyVizRowFrom(mr keyviz.MatrixRow, numCols int) *KeyVizRow {
 		Aggregate:         mr.Aggregate,
 		RouteCount:        total,
 		RouteIDsTruncated: mr.Aggregate && total > uint64(len(mr.MemberRoutes)),
-		RaftGroupID:       mr.RaftGroupID,
-		LeaderTerm:        mr.LeaderTerm,
 		Values:            make([]uint64, numCols),
+		// Per-cell parallel arrays — populated in the column-loop
+		// above, not here, because the first MatrixRow seen for a
+		// route only carries that one column's identity.
+		RaftGroupIDs: make([]uint64, numCols),
+		LeaderTerms:  make([]uint64, numCols),
 	}
 	if mr.Aggregate {
 		row.RouteIDs = append([]uint64(nil), mr.MemberRoutes...)

internal/admin/keyviz_handler_test.go

@@ -499,13 +499,17 @@ func TestKeyVizHandlerSkipsFanoutForPeerCall(t *testing.T) {
 		"recursion guard violated: handler dialled a peer despite X-Admin-Fanout-Peer being set")
 }
 
-// TestKeyVizHandlerStampsRaftIdentity pins the Phase 2-C+ wire
-// extension on the JSON path: MatrixRow.RaftGroupID and
-// MatrixRow.LeaderTerm propagate through newKeyVizRowFrom into the
-// JSON KeyVizRow's raft_group_id and leader_term fields.
+// TestKeyVizHandlerStampsRaftIdentity pins the Phase 2-C+ per-cell
+// wire extension on the JSON path: MatrixRow.RaftGroupID and
+// MatrixRow.LeaderTerm propagate per-column through the pivot into
+// the JSON KeyVizRow's raft_group_ids[] and leader_terms[] arrays
+// (parallel to values[]). A row that appears in only some columns
+// reports zero (the documented "term not tracked" sentinel) for
+// the columns where it's absent.
 func TestKeyVizHandlerStampsRaftIdentity(t *testing.T) {
 	t.Parallel()
 	t0 := time.Unix(1_700_000_000, 0)
+	t1 := t0.Add(time.Minute)
 	srv := newKeyVizTestServer(t, &fakeKeyVizSource{cols: []keyviz.MatrixColumn{
 		{
 			At: t0,
@@ -514,6 +518,13 @@ func TestKeyVizHandlerStampsRaftIdentity(t *testing.T) {
 				{RouteID: 2, Start: []byte("b"), End: []byte("c"), Writes: 9},
 			},
 		},
+		// col 1: route 1 still active (term flipped to 43); route 2 absent.
+		{
+			At: t1,
+			Rows: []keyviz.MatrixRow{
+				{RouteID: 1, Start: []byte("a"), End: []byte("b"), Writes: 11, RaftGroupID: 7, LeaderTerm: 43},
+			},
+		},
 	}})
 	defer srv.Close()
 
@@ -525,10 +536,12 @@ func TestKeyVizHandlerStampsRaftIdentity(t *testing.T) {
 	require.Len(t, matrix.Rows, 2)
 	r1, r2 := matrix.Rows[0], matrix.Rows[1]
 	require.Equal(t, "route:1", r1.BucketID)
-	require.Equal(t, uint64(7), r1.RaftGroupID, "RaftGroupID must propagate to JSON raft_group_id")
-	require.Equal(t, uint64(42), r1.LeaderTerm, "LeaderTerm must propagate to JSON leader_term")
-	// route:2 — zero values stay zero.
+	// Per-cell: col0 was term 42, col1 was term 43 (mid-window flip).
+	require.Equal(t, []uint64{7, 7}, r1.RaftGroupIDs, "raft_group_ids must be parallel to values")
+	require.Equal(t, []uint64{42, 43}, r1.LeaderTerms,
+		"leader_terms must capture per-column term so the aggregator can dedupe across the flip")
+	// route:2 — present only in col0; col1 zero is the "term not tracked" sentinel.
 	require.Equal(t, "route:2", r2.BucketID)
-	require.Equal(t, uint64(0), r2.RaftGroupID)
-	require.Equal(t, uint64(0), r2.LeaderTerm)
+	require.Equal(t, []uint64{0, 0}, r2.RaftGroupIDs)
+	require.Equal(t, []uint64{0, 0}, r2.LeaderTerms)
 }

main.go

@@ -1559,10 +1559,22 @@ type groupTermSnapshot struct {
 func publishLeaderTerms(s *keyviz.MemSampler, runtimes []*raftGroupRuntime) {
 	snaps := make([]groupTermSnapshot, 0, len(runtimes))
 	for _, rt := range runtimes {
-		if rt == nil || rt.engine == nil {
+		if rt == nil {
+			continue
+		}
+		// Snapshot rt.engine into a local once so a concurrent
+		// rt.Close() (which sets rt.engine = nil) cannot race the
+		// nil-check and the Status() call. On startup-error paths
+		// that return before joining all goroutines (e.g.
+		// setupAdminService failing), this goroutine can otherwise
+		// observe rt.engine going nil between the two reads and
+		// panic / trigger a race-detector failure. (Codex P2 on
+		// PR #720.)
+		engine := rt.engine
+		if engine == nil {
 			continue
 		}
-		snaps = append(snaps, groupTermSnapshot{groupID: rt.spec.id, term: rt.engine.Status().Term})
+		snaps = append(snaps, groupTermSnapshot{groupID: rt.spec.id, term: engine.Status().Term})
 	}
 	publishLeaderTermsFromSnapshots(s, snaps)
 }