Make restore-token reset atomic

Ultraworked with [omo](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: sisyphus-dev-ai <sisyphus-dev-ai@users.noreply.github.com>

Author: minpeter

internal/airplay/credentials.go

@@ -103,6 +103,16 @@ type CredentialStore struct {
 	backend CredentialBackend
 }
 
+// RestoreTokenReset is a one-shot restore-token deletion which can be rolled
+// back if the daemon loses its stream reservation before reconnecting.
+type RestoreTokenReset struct {
+	store         *CredentialStore
+	deviceID      string
+	previousToken string
+	changed       bool
+	done          bool
+}
+
 // NewCredentialStore creates a credential store backed by a JSON file at path.
 func NewCredentialStore(path string) (*CredentialStore, error) {
 	fb, err := newFileBackend(path)
@@ -196,19 +206,76 @@ func (cs *CredentialStore) SaveRestoreToken(deviceID, restoreToken string) error
 // ClearRestoreToken removes only the Wayland screencast restore token for a
 // device. Pairing credentials and all other device entries are preserved.
 func (cs *CredentialStore) ClearRestoreToken(deviceID string) error {
+	reset, err := cs.BeginRestoreTokenReset(deviceID)
+	if err != nil {
+		return err
+	}
+	reset.Commit()
+	return nil
+}
+
+// BeginRestoreTokenReset clears the restore token while retaining enough state
+// to restore it if the caller cannot commit the corresponding reconnect.
+func (cs *CredentialStore) BeginRestoreTokenReset(deviceID string) (*RestoreTokenReset, error) {
 	cs.mu.Lock()
 	defer cs.mu.Unlock()
 
+	reset := &RestoreTokenReset{store: cs, deviceID: deviceID}
 	creds, err := cs.backend.Lookup(deviceID)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	if creds == nil || creds.RestoreToken == "" {
-		return nil
+		return reset, nil
 	}
+	reset.previousToken = creds.RestoreToken
+	reset.changed = true
 	updated := *creds
 	updated.RestoreToken = ""
-	return cs.backend.Save(deviceID, &updated)
+	if err := cs.backend.Save(deviceID, &updated); err != nil {
+		return nil, err
+	}
+	return reset, nil
+}
+
+// Commit makes a successful deletion permanent.
+func (r *RestoreTokenReset) Commit() {
+	if r == nil {
+		return
+	}
+	r.store.mu.Lock()
+	r.done = true
+	r.store.mu.Unlock()
+}
+
+// Rollback restores only the previous token onto a fresh credential snapshot.
+// A concurrently saved non-empty token wins, and all other fields are retained.
+func (r *RestoreTokenReset) Rollback() error {
+	if r == nil {
+		return nil
+	}
+	r.store.mu.Lock()
+	defer r.store.mu.Unlock()
+	if r.done || !r.changed {
+		r.done = true
+		return nil
+	}
+	creds, err := r.store.backend.Lookup(r.deviceID)
+	if err != nil {
+		return err
+	}
+	if creds == nil {
+		creds = &SavedCredentials{}
+	}
+	if creds.RestoreToken == "" {
+		updated := *creds
+		updated.RestoreToken = r.previousToken
+		if err := r.store.backend.Save(r.deviceID, &updated); err != nil {
+			return err
+		}
+	}
+	r.done = true
+	return nil
 }
 
 // fileBackend stores credentials as a JSON file on disk.

internal/airplay/credentials_clear_test.go

@@ -111,3 +111,47 @@ func TestCredentialStoreClearRestoreTokenUsesBackendWithoutDeletingEntry(t *test
 		t.Fatalf("backend credentials after clear = %+v", creds)
 	}
 }
+
+func TestRestoreTokenResetRollbackMergesConcurrentCredentialChanges(t *testing.T) {
+	backend := &recordingCredentialBackend{devices: map[string]*SavedCredentials{
+		"device-1": {PairingID: "pair-1", RestoreToken: "restore-1"},
+	}}
+	store := NewCredentialStoreWithBackend(backend)
+
+	reset, err := store.BeginRestoreTokenReset("device-1")
+	if err != nil {
+		t.Fatalf("BeginRestoreTokenReset: %v", err)
+	}
+	backend.devices["device-1"] = &SavedCredentials{
+		PairingID:    "pair-2",
+		RestoreToken: "",
+	}
+
+	if err := reset.Rollback(); err != nil {
+		t.Fatalf("Rollback: %v", err)
+	}
+	creds := backend.devices["device-1"]
+	if creds.PairingID != "pair-2" || creds.RestoreToken != "restore-1" {
+		t.Fatalf("credentials after rollback = %+v", creds)
+	}
+}
+
+func TestRestoreTokenResetRollbackDoesNotOverwriteConcurrentToken(t *testing.T) {
+	backend := &recordingCredentialBackend{devices: map[string]*SavedCredentials{
+		"device-1": {RestoreToken: "restore-1"},
+	}}
+	store := NewCredentialStoreWithBackend(backend)
+
+	reset, err := store.BeginRestoreTokenReset("device-1")
+	if err != nil {
+		t.Fatalf("BeginRestoreTokenReset: %v", err)
+	}
+	backend.devices["device-1"] = &SavedCredentials{RestoreToken: "restore-2"}
+
+	if err := reset.Rollback(); err != nil {
+		t.Fatalf("Rollback: %v", err)
+	}
+	if got := backend.devices["device-1"].RestoreToken; got != "restore-2" {
+		t.Fatalf("restore token after rollback = %q, want concurrent token", got)
+	}
+}

internal/daemon/daemon.go

@@ -1354,24 +1354,34 @@ func (d *Daemon) getOrStartPreparedCaptureGroup(ctx context.Context, entry *acti
 	if d.captureGroups == nil {
 		d.captureGroups = make(map[videoCaptureKey]*videoCaptureGroup)
 	}
-	if group := d.captureGroups[key]; group != nil {
+	group := d.captureGroups[key]
+	var captureCtx context.Context
+	var captureCancel context.CancelFunc
+	if group != nil {
 		if group.resetReservedBy != nil && group.resetReservedBy != entry {
 			d.mu.Unlock()
 			return nil, 0, fmt.Errorf("%w: %dx%d", errCaptureGroupResetReserved, key.maxWidth, key.maxHeight)
 		}
-		entry.captureGroup = group
-		broadcast := group.broadcast
-		d.mu.Unlock()
-		preparation.Close()
-		if broadcast == nil {
-			return nil, 0, fmt.Errorf("capture group %dx%d has no broadcast", key.maxWidth, key.maxHeight)
+		if group.resetReservedBy == entry && group.broadcast == nil && group.capture == nil && group.cancel == nil {
+			captureCtx, captureCancel = context.WithCancel(context.Background())
+			group.cancel = captureCancel
+			group.resetReservedBy = nil
+		} else {
+			entry.captureGroup = group
+			broadcast := group.broadcast
+			d.mu.Unlock()
+			preparation.Close()
+			if broadcast == nil {
+				return nil, 0, fmt.Errorf("capture group %dx%d has no broadcast", key.maxWidth, key.maxHeight)
+			}
+			return broadcast, group.minimumVideoLead, nil
 		}
-		return broadcast, group.minimumVideoLead, nil
+	} else {
+		captureCtx, captureCancel = context.WithCancel(context.Background())
+		group = &videoCaptureGroup{key: key, cancel: captureCancel}
+		d.captureGroups[key] = group
+		entry.captureGroup = group
 	}
-
-	captureCtx, captureCancel := context.WithCancel(context.Background())
-	group := &videoCaptureGroup{key: key, cancel: captureCancel}
-	d.captureGroups[key] = group
 	entry.captureGroup = group
 	d.mu.Unlock()
 
@@ -1464,26 +1474,33 @@ func (d *Daemon) getOrStartCaptureGroup(entry *activeStream, restoreToken, devic
 	if d.captureGroups == nil {
 		d.captureGroups = make(map[videoCaptureKey]*videoCaptureGroup)
 	}
-	if group := d.captureGroups[key]; group != nil {
+	group := d.captureGroups[key]
+	var captureCtx context.Context
+	var captureCancel context.CancelFunc
+	if group != nil {
 		if group.resetReservedBy != nil && group.resetReservedBy != entry {
 			d.mu.Unlock()
 			return nil, fmt.Errorf("%w: %dx%d", errCaptureGroupResetReserved, key.maxWidth, key.maxHeight)
 		}
-		entry.captureGroup = group
-		broadcast := group.broadcast
-		d.mu.Unlock()
-		if broadcast == nil {
-			return nil, fmt.Errorf("capture group %dx%d has no broadcast", key.maxWidth, key.maxHeight)
+		if group.resetReservedBy == entry && group.broadcast == nil && group.capture == nil && group.cancel == nil {
+			captureCtx, captureCancel = context.WithCancel(context.Background())
+			group.cancel = captureCancel
+			group.resetReservedBy = nil
+		} else {
+			entry.captureGroup = group
+			broadcast := group.broadcast
+			d.mu.Unlock()
+			if broadcast == nil {
+				return nil, fmt.Errorf("capture group %dx%d has no broadcast", key.maxWidth, key.maxHeight)
+			}
+			return broadcast, nil
 		}
-		return broadcast, nil
+	} else {
+		captureCtx, captureCancel = context.WithCancel(context.Background())
+		group = &videoCaptureGroup{key: key, cancel: captureCancel}
+		d.captureGroups[key] = group
+		entry.captureGroup = group
 	}
-
-	// Publish the group and cancellation hook before entering the display portal
-	// or launching GStreamer. A targeted disconnect can then cancel an orphaned
-	// startup without affecting captures used by other canvas groups.
-	captureCtx, captureCancel := context.WithCancel(context.Background())
-	group := &videoCaptureGroup{key: key, cancel: captureCancel}
-	d.captureGroups[key] = group
 	entry.captureGroup = group
 	d.mu.Unlock()
 

internal/daemon/reset_restore_token.go

@@ -58,25 +58,35 @@ func (d *Daemon) handleResetRestoreToken(req Request) Response {
 	d.streamWorkers.Add(1)
 	d.mu.Unlock()
 
-	clearErr := d.credStore.ClearRestoreToken(deviceID)
+	credentialReset, clearErr := d.credStore.BeginRestoreTokenReset(deviceID)
 
 	d.mu.Lock()
 	reservationCurrent := d.restoreTokenResetReservationCurrentLocked(target, deviceID, port, entry, group)
-	if group.resetReservedBy == entry {
-		group.resetReservedBy = nil
-	}
 	state := d.overallStateLocked()
 	if d.shuttingDown {
+		if group.resetReservedBy == entry {
+			group.resetReservedBy = nil
+		}
 		d.mu.Unlock()
+		if clearErr == nil {
+			_ = credentialReset.Rollback()
+		}
 		d.streamWorkers.Done()
 		return Response{OK: false, State: state, Error: "daemon is shutting down"}
 	}
 	if !reservationCurrent {
+		if group.resetReservedBy == entry {
+			group.resetReservedBy = nil
+		}
 		d.mu.Unlock()
+		if clearErr == nil {
+			_ = credentialReset.Rollback()
+		}
 		d.streamWorkers.Done()
 		return Response{OK: false, State: state, Error: "restore token reset was canceled for " + target}
 	}
 	if clearErr != nil {
+		group.resetReservedBy = nil
 		d.mu.Unlock()
 		d.streamWorkers.Done()
 		return Response{OK: false, State: state, Error: "clear restore token: " + clearErr.Error()}
@@ -91,7 +101,17 @@ func (d *Daemon) handleResetRestoreToken(req Request) Response {
 		cancelFn:     cancel,
 		credentialCh: make(chan string, 1),
 	}
+	// Replace the old generation with an owner-only claim before releasing the
+	// daemon lock. Peers cannot create or join this key while physical cleanup
+	// runs; the designated replacement later converts the claim into a capture.
+	group.resetReservedBy = nil
 	cleanup := d.detachStreamLocked(target)
+	reservation := &videoCaptureGroup{
+		key:             group.key,
+		resetReservedBy: replacement,
+	}
+	replacement.captureGroup = reservation
+	d.captureGroups[group.key] = reservation
 	d.streams[target] = replacement
 	d.mu.Unlock()
 
@@ -112,6 +132,7 @@ func (d *Daemon) handleResetRestoreToken(req Request) Response {
 		d.mu.Unlock()
 		abandoned.run()
 		cancel()
+		_ = credentialReset.Rollback()
 		d.streamWorkers.Done()
 		if shuttingDown {
 			return Response{OK: false, State: state, Error: "daemon is shutting down"}
@@ -122,6 +143,7 @@ func (d *Daemon) handleResetRestoreToken(req Request) Response {
 	state = d.overallStateLocked()
 	d.mu.Unlock()
 
+	credentialReset.Commit()
 	go func() {
 		defer d.streamWorkers.Done()
 		d.connectAndStream(connCtx, replacement, target, port, "")