From cec20461f05b1911d3dee0d2a43dcbc8db082915 Mon Sep 17 00:00:00 2001
From: Jason Woltje <jason.woltje@uscllc.com>
Date: Fri, 7 Aug 2026 12:59:51 -0500
Subject: [PATCH] test(lease): distinguish pi probe timeouts

---
 .../src/lease-broker/invariant_r_unittest.py  | 88 +++++++++++++++++--
 1 file changed, 81 insertions(+), 7 deletions(-)

diff --git a/packages/mosaic/src/lease-broker/invariant_r_unittest.py b/packages/mosaic/src/lease-broker/invariant_r_unittest.py
index 938c4d70..2094e55c 100644
--- a/packages/mosaic/src/lease-broker/invariant_r_unittest.py
+++ b/packages/mosaic/src/lease-broker/invariant_r_unittest.py
@@ -18,6 +18,7 @@ import shutil
 import subprocess
 import sys
 import tempfile
+import time
 import unittest
 from pathlib import Path
 from typing import Final
@@ -43,9 +44,12 @@ CLAUDE_PROVEN_READ_ONLY_TOOLS: Final = frozenset({"Read", "Grep", "Glob"})
 KNOWN_DEAD_CLAUDE: Final = frozenset({"Ls", "Find"})
 
 # W-B measured Pi 0.84.1 through getAllTools(), observed every tool_call name,
-# and cross-checked dist/core/tools/index.js:17. Keep the mutating names here so
+# and cross-checked dist/core/tools/index.js:18. Keep the mutating names here so
 # a runtime registry change forces the security classification to be revisited.
 PI_VERSION: Final = "0.84.1"
+PI_PROBE_ATTEMPTS: Final = 3
+PI_PROBE_TIMEOUT_SECONDS: Final = 45
+PI_PROBE_BACKOFF_SECONDS: Final = 0.25
 PI_PROVEN_READ_ONLY_TOOLS: Final = frozenset({"read", "grep", "find", "ls"})
 PI_MUTATING_TOOLS: Final = frozenset({"bash", "edit", "write"})
 PI_MEASURED_BUILTINS: Final = PI_PROVEN_READ_ONLY_TOOLS | PI_MUTATING_TOOLS
@@ -57,6 +61,37 @@ CLAUDE_EXTRA_TOOL_ENV: Final = "MOSAIC_INVARIANT_R_CLAUDE_EXTRA_TOOL"
 PI_EXTRA_EXTENSION_ENV: Final = "MOSAIC_INVARIANT_R_PI_EXTRA_EXTENSION"
 
 
+def run_pi_registry_command(
+    command: list[str],
+    environ: dict[str, str],
+    *,
+    runner=subprocess.run,
+    sleeper=time.sleep,
+) -> subprocess.CompletedProcess[str]:
+    """Run the registry probe with bounded retries for concurrent-Pi stalls."""
+
+    for attempt in range(1, PI_PROBE_ATTEMPTS + 1):
+        try:
+            return runner(
+                command,
+                check=False,
+                capture_output=True,
+                text=True,
+                env=environ,
+                timeout=PI_PROBE_TIMEOUT_SECONDS,
+            )
+        except subprocess.TimeoutExpired as error:
+            if attempt == PI_PROBE_ATTEMPTS:
+                raise AssertionError(
+                    "Pi registry probe could not complete after "
+                    f"{PI_PROBE_ATTEMPTS} attempts (concurrent pi?); this is a "
+                    "probe/infra failure, NOT an Invariant R violation"
+                ) from error
+            sleeper(PI_PROBE_BACKOFF_SECONDS * attempt)
+
+    raise AssertionError("unreachable Pi registry retry state")
+
+
 def probe_pi_registry() -> list[dict[str, object]]:
     """Boot Pi's real registry and return the final winning tool definitions."""
 
@@ -114,13 +149,9 @@ def probe_pi_registry() -> list[dict[str, object]]:
             command.extend(("-e", extra_extension))
         command.append("Invariant R registry probe")
 
-        completed = subprocess.run(
+        completed = run_pi_registry_command(
             command,
-            check=False,
-            capture_output=True,
-            text=True,
-            env={**os.environ, "PI_OFFLINE": "1"},
-            timeout=20,
+            {**os.environ, "PI_OFFLINE": "1"},
         )
         if completed.returncode != 0 or not output.is_file():
             raise AssertionError(
@@ -200,6 +231,49 @@ class InvariantRTest(unittest.TestCase):
                     )
                     self.assertEqual(source.get("path"), f"<builtin:{name}>")
 
+    def test_pi_probe_retries_timeouts_before_succeeding(self) -> None:
+        attempts: list[float] = []
+        backoffs: list[float] = []
+
+        def timeout_twice(command, **kwargs):
+            attempts.append(kwargs["timeout"])
+            if len(attempts) < 3:
+                raise subprocess.TimeoutExpired(command, kwargs["timeout"])
+            return subprocess.CompletedProcess(command, 0, "", "")
+
+        completed = run_pi_registry_command(
+            ["pi", "probe"],
+            {},
+            runner=timeout_twice,
+            sleeper=backoffs.append,
+        )
+
+        self.assertEqual(completed.returncode, 0)
+        self.assertEqual(attempts, [45, 45, 45])
+        self.assertEqual(backoffs, [0.25, 0.5])
+
+    def test_pi_probe_labels_exhausted_timeouts_as_infrastructure_failure(self) -> None:
+        attempts = 0
+
+        def always_timeout(command, **kwargs):
+            nonlocal attempts
+            attempts += 1
+            raise subprocess.TimeoutExpired(command, kwargs["timeout"])
+
+        with self.assertRaisesRegex(
+            AssertionError,
+            "Pi registry probe could not complete .* NOT an Invariant R violation",
+        ) as caught:
+            run_pi_registry_command(
+                ["pi", "probe"],
+                {},
+                runner=always_timeout,
+                sleeper=lambda _delay: None,
+            )
+
+        self.assertEqual(attempts, 3)
+        self.assertIsInstance(caught.exception.__cause__, subprocess.TimeoutExpired)
+
 
 if __name__ == "__main__":
     unittest.main()