From 08f2c78d92bfc6323e0405b2b13a1e2351fbb6a4 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Mon, 7 Sep 2026 02:12:53 -0700
Subject: [PATCH] fix(update): cover catch-up restart clients with the unit
 budget

---
 evals/update_unit_client_budget.py            | 27 +++++++++++++++----
 hermes_cli/update_cmd_fleet.py                |  2 +-
 .../test_update_unit_client_budget.py         | 16 ++++++++---
 website/docs/developer-guide/cli-internals.md |  3 ++-
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/evals/update_unit_client_budget.py b/evals/update_unit_client_budget.py
index 1e763e0220..5297b154c1 100644
--- a/evals/update_unit_client_budget.py
+++ b/evals/update_unit_client_budget.py
@@ -7,7 +7,8 @@ import sys
 import tempfile
 import time
 
-repo, tag, output = sys.argv[1:]
+repo, tag, output = sys.argv[1:4]
+catchup = sys.argv[4:] == ["--catchup"]
 if not tag.isalnum():
     raise SystemExit("tag must be alphanumeric")
 home = Path(tempfile.mkdtemp(prefix=f"updater-{tag}-"))
@@ -15,7 +16,18 @@ allowed = {key: os.environ[key] for key in ("PATH", "XDG_RUNTIME_DIR", "DBUS_SES
 os.environ.clear()
 os.environ.update(allowed, HOME=str(home), HERMES_HOME=str(home / "hermes"))
 sys.path.insert(0, repo)
-from hermes_cli.update_cmd_fleet import _systemctl_reset_and_restart
+from hermes_cli import update_cmd_fleet as fleet
+_systemctl_reset_and_restart = fleet._systemctl_reset_and_restart
+actual_systemctl = fleet._systemctl
+def scoped_systemctl(argv, *, timeout):
+    if "list-units" in argv:
+        if "--user" not in argv:
+            return subprocess.CompletedProcess(argv, 0, "", "")
+        argv = [arg for arg in argv if arg not in ("hermes-gateway*", "hermes-serve*")] + [unit + ".service"]
+    return actual_systemctl(argv, timeout=timeout)
+if catchup:
+    # Bound discovery to our transient unit; never enumerate user services.
+    fleet._systemctl = scoped_systemctl
 unit = f"hermes-serve-audit-089c35aa-{tag}"
 cmd = ["systemctl", "--user"]
 def run(args):
@@ -24,13 +36,18 @@ def pid():
     return run(cmd + ["show", unit, "--property=MainPID", "--value"]).stdout.strip()
 record: dict[str, object] = {"repo": repo, "tag": tag, "unit": unit, "home": str(home), "tier": "native disposable systemd service"}
 try:
-    start = run(["systemd-run", "--user", "--unit", unit, "--property=ExecStop=/bin/sleep 16", "--property=TimeoutStopSec=30", "--property=TimeoutStartSec=30", "/bin/sleep", "infinity"])
+    start = run(["systemd-run", "--user", "--unit", unit, f"--property=ExecStop=/bin/sleep {31 if catchup else 16}", "--property=TimeoutStopSec=45", "--property=TimeoutStartSec=30", "/bin/sleep", "infinity"])
     assert start.returncode == 0, start.stderr
     old = pid()
     began = time.monotonic()
     try:
-        result = _systemctl_reset_and_restart(cmd, unit)
-        record.update(returncode=result.returncode, stderr=result.stderr)
+        if catchup:
+            failed = []
+            fleet._restart_systemd_gateway_units_best_effort(failed)
+            record.update(failed_units=failed)
+        else:
+            result = _systemctl_reset_and_restart(cmd, unit)
+            record.update(returncode=result.returncode, stderr=result.stderr)
     except subprocess.TimeoutExpired as exc:
         record.update(timeout=exc.timeout)
     record["elapsed"] = time.monotonic() - began
diff --git a/hermes_cli/update_cmd_fleet.py b/hermes_cli/update_cmd_fleet.py
index 1a2ba0ab45..1896391ca4 100644
--- a/hermes_cli/update_cmd_fleet.py
+++ b/hermes_cli/update_cmd_fleet.py
@@ -208,7 +208,7 @@ def _restart_systemd_gateway_units_best_effort(failed: list, listings) -> None:
             manage_cmd = list(_cmd) + ["--no-ask-password"]
             if _needs_sudo(_scope):
                 manage_cmd = ["sudo", "-n"] + manage_cmd
-            result = _systemctl_reset_and_restart(manage_cmd, svc_name)
+            result = _systemctl_reset_and_restart(manage_cmd, svc_name, scope_cmd=_cmd)
             if result.returncode != 0 or not _wait_for_service_active(_cmd, svc_name):
                 failed.append(svc_name)
 
diff --git a/tests/hermes_cli/test_update_unit_client_budget.py b/tests/hermes_cli/test_update_unit_client_budget.py
index 9d540d1a38..ef61c1953e 100644
--- a/tests/hermes_cli/test_update_unit_client_budget.py
+++ b/tests/hermes_cli/test_update_unit_client_budget.py
@@ -6,14 +6,17 @@ import pytest
 from hermes_cli import update_cmd_fleet as fleet
 
 
-@pytest.mark.parametrize("graceful,retry", [(False, False), (False, True), (True, False)])
+@pytest.mark.parametrize("graceful,retry", [(False, False), (False, True), (True, False), ("catchup", False)])
 def test_unit_transaction_budget_preserves_scope_and_health(monkeypatch, graceful, retry):
-    scope = ["systemctl", "--no-ask-password"]
-    manage = ["sudo", "-n", *scope]
+    catchup = graceful == "catchup"
+    scope = ["systemctl", "--user"] if catchup else ["systemctl", "--no-ask-password"]
+    manage = scope if catchup else ["sudo", "-n", *scope]
     calls = []
 
     def systemctl(cmd, *, timeout):
         calls.append((cmd, timeout))
+        if "list-units" in cmd:
+            return subprocess.CompletedProcess(cmd, 0, "hermes-serve-test.service loaded active running", "")
         if "show" in cmd:
             assert cmd[:len(scope)] == scope
             output = "42" if "--property=MainPID" in cmd else "TimeoutStopUSec=70s\nTimeoutStartUSec=90s"
@@ -24,6 +27,13 @@ def test_unit_transaction_budget_preserves_scope_and_health(monkeypatch, gracefu
         return subprocess.CompletedProcess(cmd, 0, "active", "")
 
     monkeypatch.setattr(fleet, "_systemctl", systemctl)
+    if catchup:
+        monkeypatch.setattr(fleet, "_SYSTEMD_SCOPES", (("user", scope),))
+        failed = []
+        fleet._restart_systemd_gateway_units_best_effort(failed)
+        assert not failed
+        assert sum("restart" in cmd for cmd, _ in calls) == 1
+        return
     monkeypatch.setattr(fleet, "_drain_or_signal_gateway_for_update", lambda *a: True)
     health = iter([False, True] if retry else [True])
     monkeypatch.setattr(fleet, "_wait_for_service_active", lambda *a, **kw: next(health))
diff --git a/website/docs/developer-guide/cli-internals.md b/website/docs/developer-guide/cli-internals.md
index 1f7b6469d8..58792e4e51 100644
--- a/website/docs/developer-guide/cli-internals.md
+++ b/website/docs/developer-guide/cli-internals.md
@@ -17,7 +17,8 @@ field failure each stage guards are documented in `hermes_cli/AGENTS.md`; user-f
 The systemd blunt-restart fallback waits for the unit's `TimeoutStopUSec` plus
 `TimeoutStartUSec`, with 15 seconds of client-side slack. It reads the target unit
 in the same manager scope as the restart; both the initial attempt and retry use
-this budget. A start after a graceful drain uses only the start budget plus slack.
+this budget, including the catch-up restart after an interrupted update.
+A start after a graceful drain uses only the start budget plus slack.
 A missing, unparseable, or infinite phase limit falls back to 90 seconds
 for that phase, keeping unattended updates bounded. Timing out the `systemctl`
 client does **not** cancel the manager's transaction. Custom multi-command stop