fix(postgres): recheck full reconciliations before witnessing them

Author: rodboev

internal/postgres/sync.go

@@ -173,6 +173,8 @@ type Sync struct {
 	// vectorSource, when set, supplies the local vectors.db active generation
 	// pushed as a phase at the end of Push. Nil disables the phase.
 	vectorSource VectorPushSource
+	// afterVectorApply is a full/scoped post-apply test hook.
+	afterVectorApply func()
 	// afterVectorGenerationLookup is a scoped-promotion test hook.
 	afterVectorGenerationLookup func()
 	// afterScopedVectorApply is a scoped-retry test hook.

internal/postgres/vector_push.go

@@ -10,6 +10,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"time"
 )
 
 // VectorGenerationInfo identifies the local embedding generation being pushed.
@@ -131,6 +132,7 @@ const vectorProgressStride = 2000
 // reset — is what survives id reuse safely.
 type vectorGeneration struct {
 	id              int64
+	createdAt       time.Time
 	halfvecType     string
 	machineRecorded bool
 }
@@ -403,6 +405,11 @@ func (s *Sync) pushVectors(
 	); err != nil {
 		return res, err
 	}
+	if s.afterVectorApply != nil {
+		hook := s.afterVectorApply
+		s.afterVectorApply = nil
+		hook()
+	}
 	if scope != nil {
 		if s.afterScopedVectorApply != nil {
 			hook := s.afterScopedVectorApply
@@ -439,6 +446,37 @@ func (s *Sync) pushVectors(
 				resolved.id, finalProbe.id,
 			)
 		}
+	} else {
+		retryGenerationWide := func(msg string, args ...any) (VectorPushResult, error) {
+			log.Printf(msg, args...)
+			_ = export.Close()
+			export = nil
+			return s.pushVectors(
+				ctx, full, nil, lastReconciledGeneration, failedSessions, onProgress,
+			)
+		}
+		finalProbe, found, err := s.lookupVectorGeneration(ctx, gen.Fingerprint, witnessKey)
+		if err != nil {
+			return res, err
+		}
+		if !found {
+			return retryGenerationWide(
+				"vector push: generation %q disappeared before recording the generation-wide witness; retrying generation-wide",
+				gen.Fingerprint,
+			)
+		}
+		if finalProbe.id != resolved.id {
+			return retryGenerationWide(
+				"vector push: active generation id changed from %d to %d before recording the generation-wide witness; retrying generation-wide",
+				resolved.id, finalProbe.id,
+			)
+		}
+		if !finalProbe.createdAt.Equal(resolved.createdAt) {
+			return retryGenerationWide(
+				"vector push: generation %d was recreated before recording the generation-wide witness; retrying generation-wide",
+				resolved.id,
+			)
+		}
 	}
 	if scope == nil && res.SessionsDeferred == 0 {
 		if err := s.recordVectorGenerationMachine(ctx, resolved.id, witnessKey); err != nil {
@@ -454,13 +492,21 @@ func (s *Sync) pushVectors(
 func (s *Sync) lookupVectorGeneration(
 	ctx context.Context, fingerprint, witnessKey string,
 ) (vectorGeneration, bool, error) {
-	genID, _, ok, err := LookupVectorGeneration(ctx, s.pg, fingerprint)
-	if err != nil {
-		return vectorGeneration{}, false, err
+	var genID int64
+	var createdAt time.Time
+	err := s.pg.QueryRowContext(ctx,
+		`SELECT id, created_at FROM vector_generations WHERE fingerprint = $1`,
+		fingerprint,
+	).Scan(&genID, &createdAt)
+	if err == sql.ErrNoRows {
+		return vectorGeneration{}, false, nil
 	}
-	if !ok {
+	if isUndefinedTable(err) {
 		return vectorGeneration{}, false, nil
 	}
+	if err != nil {
+		return vectorGeneration{}, false, fmt.Errorf("looking up vector generation: %w", err)
+	}
 	extSchema, err := vectorExtensionSchema(ctx, s.pg)
 	if err != nil {
 		return vectorGeneration{}, false, err
@@ -477,6 +523,7 @@ SELECT EXISTS (
 	}
 	return vectorGeneration{
 		id:              genID,
+		createdAt:       createdAt,
 		halfvecType:     extSchema + ".halfvec",
 		machineRecorded: machineRecorded,
 	}, true, nil
@@ -824,6 +871,12 @@ func (s *Sync) resolveVectorGeneration(
 	if err != nil {
 		return vectorGeneration{}, err
 	}
+	var createdAt time.Time
+	if err := s.pg.QueryRowContext(ctx,
+		`SELECT created_at FROM vector_generations WHERE id = $1`, genID,
+	).Scan(&createdAt); err != nil {
+		return vectorGeneration{}, fmt.Errorf("reading vector generation created_at: %w", err)
+	}
 	var machineRecorded bool
 	if err := s.pg.QueryRowContext(ctx, `
 SELECT EXISTS (
@@ -834,6 +887,7 @@ SELECT EXISTS (
 	}
 	return vectorGeneration{
 		id:              genID,
+		createdAt:       createdAt,
 		halfvecType:     extSchema + ".halfvec",
 		machineRecorded: machineRecorded,
 	}, nil

internal/postgres/vector_push_pg_test.go

@@ -928,6 +928,90 @@ SELECT COUNT(*) FROM vector_push_state
 		"both beta sessions have state rows under the recreated generation")
 }
 
+// TestVectorPushFullRechecksGenerationBeforeRecordingWitness pins the last
+// incarnation race: a full reconciliation must verify the generation survived
+// intact before it records the durable witness that later scoped pushes trust.
+func TestVectorPushFullRechecksGenerationBeforeRecordingWitness(t *testing.T) {
+	pgURL := testPGURL(t)
+	sync, localDB, pg := newVectorPushTestSync(
+		t, pgURL, "agentsview_vector_push_full_witness_race_test")
+	ctx := context.Background()
+
+	seedVectorSession(t, localDB, "A")
+	seedVectorSession(t, localDB, "B")
+	src := &fakeVectorSource{
+		gen:    VectorGenerationInfo{Fingerprint: "fp-full-race", Model: "m", Dimension: 4},
+		hasGen: true,
+		hashes: map[string]string{"A": "a1", "B": "b1"},
+		docs: map[string][]VectorPushDoc{
+			"A": {vdoc("A", "A#0", 0, "ca1", "a1", []float32{1, 0, 0, 0})},
+			"B": {vdoc("B", "B#0", 0, "cb1", "b1", []float32{0, 1, 0, 0})},
+		},
+	}
+	sync.vectorSource = src
+
+	_, err := sync.Push(ctx, false, nil)
+	require.NoError(t, err, "baseline Push")
+
+	var gen1 int64
+	require.NoError(t, pg.QueryRow(
+		`SELECT id FROM vector_generations WHERE fingerprint = $1`, "fp-full-race",
+	).Scan(&gen1), "gen1 id")
+	witnessKey, err := sync.vectorGenerationWitnessKey()
+	require.NoError(t, err, "vectorGenerationWitnessKey")
+
+	src.genScopes = nil
+	src.hashScopes = nil
+	src.hashes = map[string]string{"A": "a2", "B": "b2"}
+	src.docs = map[string][]VectorPushDoc{
+		"A": {vdoc("A", "A#0", 0, "ca2", "a2", []float32{0, 0, 1, 0})},
+		"B": {vdoc("B", "B#0", 0, "cb2", "b2", []float32{0, 0, 0, 1})},
+	}
+	sync.afterVectorApply = func() {
+		for _, q := range []string{
+			`DROP TABLE IF EXISTS ` + vectorChunkTable(gen1),
+			`DROP TABLE IF EXISTS vector_push_state`,
+			`DROP TABLE IF EXISTS vector_generation_machines`,
+			`DROP TABLE IF EXISTS vector_documents`,
+			`DROP TABLE IF EXISTS vector_generations`,
+		} {
+			_, err := pg.Exec(q)
+			require.NoError(t, err, q)
+		}
+		unavailable, err := ensureVectorBaseSchemaPG(ctx, pg)
+		require.NoError(t, err, "recreate vector base schema")
+		require.Empty(t, unavailable, "pgvector must stay available in test")
+		gen2, err := ensureVectorGeneration(
+			ctx, pg, src.gen.Fingerprint, src.gen.Model, src.gen.Dimension,
+		)
+		require.NoError(t, err, "re-register generation")
+		require.Equal(t, gen1, gen2,
+			"recreated tables restart the id sequence, which is the case under test")
+	}
+
+	vres, err := sync.pushVectors(ctx, false, nil, gen1, nil, nil)
+	require.NoError(t, err, "full push across a pre-witness recreation")
+	assert.Equal(t, 2, vres.SessionsPushed,
+		"the recreated generation is repopulated generation-wide before success returns")
+	require.Len(t, src.genScopes, 2)
+	assert.Nil(t, src.genScopes[0])
+	assert.Nil(t, src.genScopes[1],
+		"the retry must reopen the full generation after the recreation")
+	require.Len(t, src.hashScopes, 2)
+	assert.Nil(t, src.hashScopes[0])
+	assert.Nil(t, src.hashScopes[1],
+		"both passes are generation-wide reads")
+	assert.Equal(t, gen1, vres.GenerationID,
+		"the full retry reconciles the recreated generation even when the id is reused")
+	assert.Equal(t, 1, countRows(t, pg, `
+SELECT COUNT(*) FROM vector_generation_machines
+ WHERE generation_id = $1 AND machine = $2`, gen1, witnessKey),
+		"the final witness is recorded only after the stable retry")
+	assert.Equal(t, 2, countRows(t, pg, `
+SELECT COUNT(*) FROM vector_push_state WHERE generation_id = $1`, gen1),
+		"both sessions have state rows under the recreated generation")
+}
+
 func TestVectorPushRoundTrip(t *testing.T) {
 	pgURL := testPGURL(t)
 	sync, localDB, pg := newVectorPushTestSync(