Preserve Local Changes During Fast Updates

openpilot/starpilot/software/fast_update.py

@@ -136,8 +136,6 @@ def fast_update(repo, branch, *, params, parked, expected_commit, current_os=Non
       raise FastUpdateError('Fast Update requires fresh parked state')
     if git('symbolic-ref', '--short', 'HEAD').strip() != initial_branch:
       raise FastUpdateError('Installed branch changed; refresh and try again')
-    if git('status', '--porcelain', '--untracked-files=no').strip():
-      raise FastUpdateError('Tracked source has local changes')
 
   git('check-ref-format', '--branch', branch)
   initial_branch = git('symbolic-ref', '--short', 'HEAD').strip()
@@ -178,7 +176,7 @@ def fast_update(repo, branch, *, params, parked, expected_commit, current_os=Non
   launch = git('show', f'{target}:launch_env.sh')
   if not current_os or required_os_version(launch) != current_os:
     raise FastUpdateIncompatible('Update requires a different or unverified operating system')
-  if target == previous and branch == initial_branch:
+  if target == previous and branch == initial_branch and not git('status', '--porcelain', '--untracked-files=no').strip():
     notify('complete', 'Already up to date.')
     return 'up-to-date'
   for revision in (previous, target):
@@ -218,6 +216,15 @@ def fast_update(repo, branch, *, params, parked, expected_commit, current_os=Non
 
   retained = optional('rev-parse', '--verify', 'refs/starpilot/previous^{commit}')
   retained_branch = optional('config', '--local', '--get', 'starpilot.previousVersion')
+  local_backup = None
+  if git('status', '--porcelain', '--untracked-files=no').strip():
+    notify('preparing', 'Saving local tracked changes before installation...')
+    local_backup = git('stash', 'create').strip()
+    if not re.fullmatch(r'[0-9a-f]{40}', local_backup):
+      raise FastUpdateError('Could not back up local tracked changes')
+    backup_ref = f'refs/starpilot/local-changes/{time.time_ns()}'
+    git('update-ref', backup_ref, local_backup)
+    notify('preparing', f'Local changes saved as {backup_ref}.')
   git('update-ref', 'refs/starpilot/previous', previous)
   git('config', '--local', 'starpilot.previousVersion', f'{previous} {initial_branch}')
   try:
@@ -250,6 +257,8 @@ def fast_update(repo, branch, *, params, parked, expected_commit, current_os=Non
       git('reset', '--hard', '--no-recurse-submodules', previous)
     if git('rev-parse', 'HEAD^{commit}').strip() != previous:
       raise FastUpdateError('Source rollback failed') from None
+    if local_backup is not None:
+      git('stash', 'apply', '--index', local_backup)
     if retained is None:
       git('update-ref', '-d', 'refs/starpilot/previous')
     else:

openpilot/starpilot/software/tests/test_fast_update.py

@@ -344,19 +344,67 @@ class TestFastUpdate(unittest.TestCase):
     self.assertFalse(self.invalidations)
     self.assertFalse(self.params.calls)
 
-  def test_dirty_tracked_source_rejected(self):
-    (self.repo / 'openpilot/starpilot/ui/source.py').write_text('local repair')
-    with self.assertRaises(owner.FastUpdateError):
-      self.update()
-    self.assertEqual((self.repo / 'openpilot/starpilot/ui/source.py').read_text(), 'local repair')
+  def test_dirty_tracked_source_backed_up_before_update(self):
+    path = 'openpilot/starpilot/ui/source.py'
+    (self.repo / path).write_text('local repair')
+    self.assertEqual(self.update(), 'reboot-requested')
+    self.assertEqual((self.repo / path).read_text(), 'new')
+    refs = self.git(self.repo, 'for-each-ref', '--format=%(refname)', 'refs/starpilot/local-changes').splitlines()
+    self.assertEqual(len(refs), 1)
+    self.assertEqual(self.git(self.repo, 'show', refs[0] + ':' + path), 'local repair')
 
-  def test_staged_source_rejected(self):
-    (self.repo / 'openpilot/starpilot/ui/source.py').write_text('staged repair')
-    self.git(self.repo, 'add', 'openpilot/starpilot/ui/source.py')
-    with self.assertRaises(owner.FastUpdateError):
-      self.update()
-    self.assertEqual((self.repo / 'openpilot/starpilot/ui/source.py').read_text(), 'staged repair')
+  def test_staged_and_unstaged_source_restored_on_failure(self):
+    path = 'openpilot/starpilot/ui/source.py'
+    (self.repo / path).write_text('staged repair')
+    self.git(self.repo, 'add', path)
+    (self.repo / path).write_text('unstaged repair')
+    with patch.object(self.params, 'put_bool', side_effect=RuntimeError('reboot failed')):
+      with self.assertRaisesRegex(RuntimeError, 'reboot failed'):
+        self.update()
+    self.assertEqual(self.git(self.repo, 'rev-parse', 'HEAD'), self.previous)
+    self.assertEqual(self.git(self.repo, 'show', ':' + path), 'staged repair')
+    self.assertEqual((self.repo / path).read_text(), 'unstaged repair')
+
+  def test_dirty_version_listing_leaves_changes_and_backup_refs_untouched(self):
+    path = self.repo / 'openpilot/starpilot/ui/source.py'
+    path.write_text('local repair')
+    self.assertEqual(self.update(versions_only=True), 'versions-listed')
+    self.assertEqual(path.read_text(), 'local repair')
+    self.assertEqual(self.git(self.repo, 'for-each-ref', 'refs/starpilot/local-changes'), '')
+
+  def test_backup_failure_does_not_reset_local_changes(self):
+    path = self.repo / 'openpilot/starpilot/ui/source.py'
+    path.write_text('local repair')
+    def failing(command, cwd):
+      if 'stash' in command:
+        raise RuntimeError('backup failed')
+      return owner.run(command, cwd)
+    with self.assertRaisesRegex(RuntimeError, 'backup failed'):
+      self.update(run=failing)
+    self.assertEqual(path.read_text(), 'local repair')
     self.assertFalse(self.invalidations)
+    self.assertFalse(self.params.calls)
+
+  def test_dirty_current_version_is_reinstalled_with_backup(self):
+    self.git(self.remote, 'reset', '--hard', self.previous)
+    path = self.repo / 'openpilot/starpilot/ui/source.py'
+    path.write_text('local repair')
+    self.assertEqual(self.update(), 'reboot-requested')
+    self.assertEqual(path.read_text(), 'old')
+    self.assertTrue(self.git(self.repo, 'for-each-ref', 'refs/starpilot/local-changes'))
+
+  def test_backup_ref_failure_does_not_reset_or_invalidate(self):
+    path = self.repo / 'openpilot/starpilot/ui/source.py'
+    path.write_text('local repair')
+    def failing(command, cwd):
+      if 'update-ref' in command and any('local-changes/' in item for item in command):
+        raise RuntimeError('backup ref failed')
+      return owner.run(command, cwd)
+    with self.assertRaisesRegex(RuntimeError, 'backup ref failed'):
+      self.update(run=failing)
+    self.assertEqual(path.read_text(), 'local repair')
+    self.assertFalse(self.invalidations)
+    self.assertFalse(self.params.calls)
 
   def test_reboot_write_failure_rolls_back(self):
     def failing(*args, **kwargs):