Preserve legacy ownership across machine renames

Author: rodboev

internal/postgres/collision_pgtest_test.go

@@ -84,7 +84,7 @@ func TestPushSessionGuardsAgainstCrossMachineCollision(t *testing.T) {
 	// Execute pushSession.
 	tx, err := pg.BeginTx(ctx, nil)
 	require.NoError(t, err, "BeginTx")
-	err = sync.pushSession(ctx, tx, sess, markerID)
+	err = sync.pushSession(ctx, tx, sess, markerID, "")
 	require.ErrorIs(t, err, errSessionOwnershipConflict, "pushSession should return ownership conflict sentinel")
 	require.NoError(t, tx.Commit(), "Commit")
 
@@ -156,7 +156,7 @@ func TestPushSessionAllowsMachineRenameForSameOwnerMarker(t *testing.T) {
 
 	tx, err := pg.BeginTx(ctx, nil)
 	require.NoError(t, err, "BeginTx")
-	require.NoError(t, sync.pushSession(ctx, tx, sess, markerID), "pushSession")
+	require.NoError(t, sync.pushSession(ctx, tx, sess, markerID, ""), "pushSession")
 	require.NoError(t, tx.Commit(), "Commit")
 
 	var machine, ownerMarker string
@@ -215,7 +215,7 @@ func TestPushSessionAdoptsLegacyLocalSentinelRow(t *testing.T) {
 	require.NoError(t, err, "BeginTx")
 	markerID, err := sync.pushMarkerID()
 	require.NoError(t, err, "pushMarkerID")
-	require.NoError(t, sync.pushSession(ctx, tx, sess, markerID), "pushSession")
+	require.NoError(t, sync.pushSession(ctx, tx, sess, markerID, ""), "pushSession")
 	require.NoError(t, tx.Commit(), "Commit")
 
 	var machine, ownerMarker string

internal/postgres/name_source_pgtest_test.go

@@ -62,7 +62,7 @@ func TestPushSessionNameRoundTrip(t *testing.T) {
 	// Push via pushSession directly.
 	tx, err := pg.BeginTx(ctx, nil)
 	require.NoError(t, err, "BeginTx")
-	if err := sync.pushSession(ctx, tx, sess, markerID); err != nil {
+	if err := sync.pushSession(ctx, tx, sess, markerID, ""); err != nil {
 		_ = tx.Rollback()
 		t.Fatalf("pushSession: %v", err)
 	}
@@ -105,7 +105,7 @@ func TestPushSessionNameRoundTrip(t *testing.T) {
 
 	tx2, err := pg.BeginTx(ctx, nil)
 	require.NoError(t, err, "BeginTx (second)")
-	if err := sync.pushSession(ctx, tx2, sess, markerID); err != nil {
+	if err := sync.pushSession(ctx, tx2, sess, markerID, ""); err != nil {
 		_ = tx2.Rollback()
 		t.Fatalf("pushSession (second): %v", err)
 	}

internal/postgres/push.go

@@ -85,6 +85,10 @@ func (s *Sync) Push(
 	if err != nil {
 		return result, err
 	}
+	previousMarkerMachine, markerExists, err := s.pgPushMarkerMachine(ctx, markerID)
+	if err != nil {
+		return result, err
+	}
 	if full {
 		lastPush = ""
 		// Caller requested a full push — the PG schema
@@ -113,17 +117,14 @@ func (s *Sync) Push(
 	// was reset (schema dropped, DB recreated, etc.). Force a full
 	// push so all sessions are re-synced.
 	if lastPush != "" {
-		markerExists, cErr := s.pgPushMarkerExists(ctx, markerID)
-		if cErr != nil {
-			return result, cErr
-		}
 		if !markerExists {
 			log.Printf(
 				"pgsync: local watermark set but PG push marker " +
 					"missing; PG was reset, forcing full push",
 			)
 			lastPush = ""
 			full = true
+			previousMarkerMachine = ""
 			s.schemaMu.Lock()
 			s.schemaDone = false
 			s.schemaMu.Unlock()
@@ -278,7 +279,7 @@ func (s *Sync) Push(
 		batch := sessions[i:end]
 
 		batchResult, err := s.pushBatch(
-			ctx, batch, full, markerID, &pushed,
+			ctx, batch, full, markerID, previousMarkerMachine, &pushed,
 		)
 		if err != nil {
 			return result, err
@@ -293,7 +294,7 @@ func (s *Sync) Push(
 			for _, sess := range batch {
 				sr, retryErr := s.pushBatch(
 					ctx, []db.Session{sess},
-					full, markerID, &pushed,
+					full, markerID, previousMarkerMachine, &pushed,
 				)
 				if retryErr != nil {
 					return result, retryErr
@@ -369,29 +370,33 @@ func (s *Sync) Push(
 	return result, nil
 }
 
-// pgPushMarkerExists reports whether this host's push marker is present in PG.
+// pgPushMarkerMachine reports whether this host's push marker is present in PG
+// and returns the machine value stored with the marker.
 // A missing marker while the local watermark is set means PG was reset (schema
 // dropped or recreated) since this host last pushed, so a full re-push is
 // needed. Counting rows by machine cannot detect this reliably: another host
 // pushing to the same PG can repopulate rows under a machine value this host
 // also writes -- a remote host's sessions synced in over SSH, or this host's
 // own renamed identity -- masking the loss of this host's own rows. The marker
 // is per-local-DB, so no other pusher can satisfy this check.
-func (s *Sync) pgPushMarkerExists(ctx context.Context, markerID string) (bool, error) {
-	var exists bool
+func (s *Sync) pgPushMarkerMachine(ctx context.Context, markerID string) (string, bool, error) {
+	var machine string
 	err := s.pg.QueryRowContext(ctx,
-		`SELECT EXISTS (SELECT 1 FROM sync_metadata WHERE key = $1)`,
+		`SELECT value FROM sync_metadata WHERE key = $1`,
 		pushMarkerKeyPrefix+markerID,
-	).Scan(&exists)
+	).Scan(&machine)
 	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return "", false, nil
+		}
 		if isUndefinedTable(err) {
-			return false, nil
+			return "", false, nil
 		}
-		return false, fmt.Errorf(
+		return "", false, fmt.Errorf(
 			"checking pg push marker: %w", err,
 		)
 	}
-	return exists, nil
+	return machine, true, nil
 }
 
 // writePushMarker records this host's push marker in PG so a later push can
@@ -452,6 +457,7 @@ func (s *Sync) pushBatch(
 	batch []db.Session,
 	full bool,
 	markerID string,
+	previousMarkerMachine string,
 	pushed *[]db.Session,
 ) (batchResult, error) {
 	tx, err := s.pg.BeginTx(ctx, nil)
@@ -466,7 +472,7 @@ func (s *Sync) pushBatch(
 	skippedConflicts := 0
 	for _, sess := range batch {
 		if err := s.pushSession(
-			ctx, tx, sess, markerID,
+			ctx, tx, sess, markerID, previousMarkerMachine,
 		); err != nil {
 			if errors.Is(err, errSessionOwnershipConflict) {
 				skippedConflicts++
@@ -798,7 +804,10 @@ func pushedSessionMachine(sess db.Session, fallbackMachine string) string {
 	return fallbackMachine
 }
 
-func sameSessionOwner(existingOwnerMarker, existingMachine, markerID, pushedMachine string) bool {
+func sameSessionOwner(
+	existingOwnerMarker, existingMachine, markerID, pushedMachine,
+	previousMarkerMachine string,
+) bool {
 	if existingOwnerMarker != "" {
 		return existingOwnerMarker == markerID
 	}
@@ -808,6 +817,9 @@ func sameSessionOwner(existingOwnerMarker, existingMachine, markerID, pushedMach
 	if existingMachine == "local" {
 		return true
 	}
+	if previousMarkerMachine != "" && existingMachine == previousMarkerMachine {
+		return true
+	}
 	return existingMachine == pushedMachine
 }
 
@@ -872,7 +884,7 @@ func nilStrTS(s *string) any {
 // local-only and used solely by the sync engine to detect
 // re-parsed sessions.
 func (s *Sync) pushSession(
-	ctx context.Context, tx *sql.Tx, sess db.Session, markerID string,
+	ctx context.Context, tx *sql.Tx, sess db.Session, markerID, previousMarkerMachine string,
 ) error {
 	createdAt, _ := ParseSQLiteTimestamp(sess.CreatedAt)
 	isAutomated := sess.IsAutomated
@@ -890,6 +902,7 @@ func (s *Sync) pushSession(
 		existingMachine.String,
 		markerID,
 		pushedMachine,
+		previousMarkerMachine,
 	) {
 		log.Printf(
 			"pgsync: session %s: skipping — already owned by machine %q, "+
@@ -989,7 +1002,9 @@ func (s *Sync) pushSession(
 		WHERE ((
 				sessions.owner_marker = ''
 				AND (sessions.machine = EXCLUDED.machine
-					OR sessions.machine = 'local')
+					OR sessions.machine = 'local'
+					OR sessions.machine = ''
+					OR sessions.machine = $48)
 			)
 			OR sessions.owner_marker = EXCLUDED.owner_marker)
 			AND (
@@ -1070,6 +1085,7 @@ func (s *Sync) pushSession(
 		sess.HealthScore, nilStr(sess.HealthGrade),
 		sess.HasToolCalls, sess.HasContextData,
 		sess.SecretLeakCount, sess.SecretsRulesVersion,
+		previousMarkerMachine,
 	)
 	if err != nil {
 		return err
@@ -1084,7 +1100,10 @@ func (s *Sync) pushSession(
 			currentOwnerMarker = existingOwnerMarker.String
 			currentMachine = existingMachine.String
 		}
-		if refreshErr == nil && !sameSessionOwner(currentOwnerMarker, currentMachine, markerID, pushedMachine) {
+		if refreshErr == nil && !sameSessionOwner(
+			currentOwnerMarker, currentMachine, markerID, pushedMachine,
+			previousMarkerMachine,
+		) {
 			log.Printf(
 				"pgsync: session %s: skipping — already owned by machine %q, this pusher is %q; sync from the origin machine to update",
 				sess.ID, currentMachine, pushedMachine,

internal/postgres/push_pgtest_test.go

@@ -373,7 +373,7 @@ func TestPushSessionTerminationStatus(t *testing.T) {
 		t.Helper()
 		tx, err := pg.BeginTx(ctx, nil)
 		require.NoError(t, err, "BeginTx")
-		if err := sync.pushSession(ctx, tx, s, markerID); err != nil {
+		if err := sync.pushSession(ctx, tx, s, markerID, ""); err != nil {
 			_ = tx.Rollback()
 			t.Fatalf("pushSession: %v", err)
 		}
@@ -439,7 +439,7 @@ func TestPushSessionPreservesSourceMachine(t *testing.T) {
 	require.NoError(t, err, "BeginTx")
 	markerID, err := sync.pushMarkerID()
 	require.NoError(t, err, "pushMarkerID")
-	require.NoError(t, sync.pushSession(ctx, tx, remoteSession, markerID), "pushSession")
+	require.NoError(t, sync.pushSession(ctx, tx, remoteSession, markerID, ""), "pushSession")
 	require.NoError(t, tx.Commit(), "Commit")
 
 	var got string
@@ -1101,6 +1101,87 @@ func TestPushUpdatesSentinelMachineWhenSyncMachineChanges(t *testing.T) {
 		"sentinel machine must follow the new fallback")
 }
 
+func TestPushAdoptsOwnerlessRowsFromPreviousMarkerMachine(t *testing.T) {
+	pgURL := testPGURL(t)
+
+	const schema = "agentsview_push_legacy_marker_machine_test"
+	pg, err := Open(pgURL, schema, true)
+	require.NoError(t, err, "Open")
+	defer pg.Close()
+
+	ctx := context.Background()
+	_, err = pg.Exec(`DROP SCHEMA IF EXISTS ` + schema + ` CASCADE`)
+	require.NoError(t, err, "drop schema")
+	require.NoError(t, EnsureSchema(ctx, pg, schema), "EnsureSchema")
+
+	localDB, err := db.Open(filepath.Join(t.TempDir(), "local.db"))
+	require.NoError(t, err, "db.Open")
+	defer localDB.Close()
+
+	const markerID = "legacy-marker-1"
+	require.NoError(t, localDB.SetSyncState("pg_push_marker_id", markerID),
+		"seed local push marker")
+
+	sync := &Sync{
+		pg:         pg,
+		local:      localDB,
+		machine:    "host-b",
+		schema:     schema,
+		schemaDone: true,
+	}
+
+	const sessID = "legacy-previous-machine-1"
+	require.NoError(t, localDB.UpsertSession(db.Session{
+		ID:           sessID,
+		Project:      "proj",
+		Machine:      "host-b",
+		Agent:        "claude",
+		MessageCount: 1,
+		CreatedAt:    "2026-01-01T00:00:00Z",
+	}), "UpsertSession")
+	require.NoError(t, localDB.InsertMessages([]db.Message{{
+		SessionID:     sessID,
+		Ordinal:       0,
+		Role:          "assistant",
+		Content:       "hello",
+		ContentLength: 5,
+	}}), "InsertMessages")
+
+	_, err = pg.ExecContext(ctx, `
+		INSERT INTO sync_metadata (key, value)
+		VALUES ($1, $2)
+	`, pushMarkerKeyPrefix+markerID, "host-a")
+	require.NoError(t, err, "seed previous marker machine")
+	_, err = pg.ExecContext(ctx, `
+		INSERT INTO sessions (
+			id, machine, owner_marker, project, agent, created_at
+		) VALUES ($1, $2, $3, $4, $5, NOW())
+	`, sessID, "host-a", "", "proj", "claude")
+	require.NoError(t, err, "seed ownerless legacy session")
+
+	res, err := sync.Push(ctx, true, nil)
+	require.NoError(t, err, "Push")
+	assert.Zero(t, res.Errors, "push should report no failed sessions")
+	assert.Zero(t, res.SkippedConflicts,
+		"previous-marker-machine legacy row should be adopted")
+	assert.Equal(t, 1, res.SessionsPushed,
+		"legacy row should be counted as pushed")
+
+	var machine, ownerMarker string
+	require.NoError(t, pg.QueryRow(
+		`SELECT machine, owner_marker FROM sessions WHERE id = $1`, sessID,
+	).Scan(&machine, &ownerMarker), "reading adopted row")
+	assert.Equal(t, "host-b", machine)
+	assert.Equal(t, markerID, ownerMarker)
+
+	var markerMachine string
+	require.NoError(t, pg.QueryRow(
+		`SELECT value FROM sync_metadata WHERE key = $1`,
+		pushMarkerKeyPrefix+markerID,
+	).Scan(&markerMachine), "reading marker machine")
+	assert.Equal(t, "host-b", markerMachine)
+}
+
 func TestPushReportsSkippedConflicts(t *testing.T) {
 	pgURL := testPGURL(t)