vmupdate: wait for PCI-assigned VM shutdown in update flow

Author: Jayant-kernel

vmupdate/qube_connection.py

@@ -22,12 +22,14 @@
 import shutil
 import signal
 import tempfile
+import asyncio
 import concurrent.futures
 from os.path import join
 from subprocess import CalledProcessError
 from typing import List
 
 import qubesadmin
+from qubesadmin.events.utils import wait_for_domain_shutdown
 from vmupdate.agent.source.args import AgentArgs
 from vmupdate.agent.source.log_config import LOGPATH, LOG_FILE
 from vmupdate.agent.source.status import StatusInfo, FinalStatus
@@ -94,9 +96,19 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if self.qube.is_running() and not self._initially_running:
             self.logger.info('Shutdown %s', self.qube.name)
             self.qube.shutdown()
+            if self._has_assigned_pci_devices(self.qube):
+                asyncio.run(wait_for_domain_shutdown([self.qube]))
 
         self.__connected = False
 
+    @staticmethod
+    def _has_assigned_pci_devices(vm) -> bool:
+        """Return True when VM has assigned PCI devices."""
+        try:
+            return any(vm.devices['pci'].get_assigned_devices())
+        except Exception:  # best effort check only
+            return False
+
     def transfer_agent(self, src_dir: str) -> ProcessResult:
         """
         Copy a directory content to the workdir in the qube.

vmupdate/tests/test_qube_connection.py

@@ -0,0 +1,77 @@
+# coding=utf-8
+#
+# The Qubes OS Project, https://www.qubes-os.org
+#
+# Copyright (C) 2026
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 2
+# of the License, or (at your option) any later version.
+#
+from unittest.mock import Mock, patch
+
+from vmupdate.qube_connection import QubeConnection
+
+
+@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
+@patch("vmupdate.qube_connection.asyncio.run")
+def test_wait_for_shutdown_when_vm_started_by_update(_arun, wait_shutdown):
+    vm = Mock()
+    vm.name = "hvm1"
+    vm.is_running.side_effect = [False, True]
+    vm.devices = {'pci': Mock()}
+    vm.devices['pci'].get_assigned_devices.return_value = ["00_1f.2"]
+    status_notifier = Mock()
+    logger = Mock()
+
+    with QubeConnection(
+            vm, "/tmp/qubes-update", cleanup=False, logger=logger,
+            show_progress=False, status_notifier=status_notifier):
+        pass
+
+    vm.shutdown.assert_called_once_with()
+    wait_shutdown.assert_called_once_with([vm])
+    _arun.assert_called_once()
+
+
+@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
+@patch("vmupdate.qube_connection.asyncio.run")
+def test_do_not_wait_for_shutdown_without_assigned_pci(_arun, wait_shutdown):
+    vm = Mock()
+    vm.name = "hvm2"
+    vm.is_running.side_effect = [False, True]
+    vm.devices = {'pci': Mock()}
+    vm.devices['pci'].get_assigned_devices.return_value = []
+    status_notifier = Mock()
+    logger = Mock()
+
+    with QubeConnection(
+            vm, "/tmp/qubes-update", cleanup=False, logger=logger,
+            show_progress=False, status_notifier=status_notifier):
+        pass
+
+    vm.shutdown.assert_called_once_with()
+    wait_shutdown.assert_not_called()
+    _arun.assert_not_called()
+
+
+@patch("vmupdate.qube_connection.wait_for_domain_shutdown")
+@patch("vmupdate.qube_connection.asyncio.run")
+def test_do_not_shutdown_if_vm_was_already_running(_arun, wait_shutdown):
+    vm = Mock()
+    vm.name = "hvm3"
+    vm.is_running.return_value = True
+    vm.devices = {'pci': Mock()}
+    vm.devices['pci'].get_assigned_devices.return_value = ["00_1f.2"]
+    status_notifier = Mock()
+    logger = Mock()
+
+    with QubeConnection(
+            vm, "/tmp/qubes-update", cleanup=False, logger=logger,
+            show_progress=False, status_notifier=status_notifier):
+        pass
+
+    vm.shutdown.assert_not_called()
+    wait_shutdown.assert_not_called()
+    _arun.assert_not_called()