From 49804a8774cc1532542ee649849989dece35b2d8 Mon Sep 17 00:00:00 2001 From: maxlamagna Date: Tue, 11 Aug 2026 21:48:59 +0100 Subject: [PATCH] feat(registry): add a starting state with a single activation path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Groundwork for an opt-in readiness gate. An instance may now be registered as `starting`: it receives the identity details a wrapper needs in order to launch, but it is not yet active. - `register(ready_gate=True)` enters `starting`, and `mark_ready` is the only path to active, so an instance cannot drift into active by another route. - `mark_starting` re-gates an existing starting or active instance, for restarts. - `claim` refuses a starting instance, so a second process cannot take an identity that is still coming up. - `cancel_starting` removes the instance with no reservation and no reclaimable entry, via a shared `_remove_instance` primitive, so the family rename-back bookkeeping runs on both removal paths. It returns that removal record rather than a bare bool, so a caller can migrate chat identity and history on a rename-back exactly as it already does for deregister. - `expire_crashed` is the crash sweep's removal. The two cases need opposite treatment — a starting instance is cancelled so a relaunch reacquires the bare name and its token dies, anything else stays reclaimable so a sleeping machine recovers — and that choice is state-dependent, so it is made under the same lock that performs the removal. Choosing outside it is a race: a wrapper re-gating from active to starting in the window would still be deregistered reclaimably, reserving its name and leaving a revivable token. Token recovery preserves `starting`. On startup every persisted instance is reloaded as reclaimable and a live wrapper transparently recovers its identity through resolve_token; that path previously forced every recovered instance to active, so restarting the server would have opened the gate for an agent whose CLI was never proven ready. Recovery of an active instance is unchanged. Inert without `ready_gate=True`: existing registration behaviour is unchanged, which is why this can land ahead of the server, UI and wrapper pieces that use it. The existing removals are untouched: `deregister` and `cancel_starting` still return exactly `{'ok': True}` plus an optional `_renamed_back`. The `cancelled` key belongs to the state-dependent expiry alone, and the by-state mode chooses both branches itself rather than inheriting them from the caller. 21 unit tests, written red-first, each guard verified by a mutation that removes it. Two cover the expiry race specifically: a barrier proving a concurrent re-gate cannot complete while the removal holds the lock, and one that fails any implementation which decides from a state read before the removal rather than during it. A further set pins the exact return shape of every removal, so the gate cannot quietly widen an existing contract. --- registry.py | 169 ++++++++++++++--- tests/test_ready_gate_registry.py | 302 ++++++++++++++++++++++++++++++ 2 files changed, 445 insertions(+), 26 deletions(-) create mode 100644 tests/test_ready_gate_registry.py diff --git a/registry.py b/registry.py index 79dff25b..187b1b46 100644 --- a/registry.py +++ b/registry.py @@ -27,7 +27,7 @@ class Instance: identity_id: str = field(default_factory=lambda: uuid.uuid4().hex) token: str = field(default_factory=lambda: secrets.token_hex(16)) epoch: int = 1 - state: str = "pending" # "pending" | "active" + state: str = "pending" # "pending" | "starting" | "active" registered_at: float = field(default_factory=time.time) @@ -183,11 +183,17 @@ def _evict_reclaimable_collisions_locked(self): # --- Registration --- - def register(self, base: str, label: str | None = None) -> dict | None: + def register(self, base: str, label: str | None = None, + ready_gate: bool = False) -> dict | None: """Register a new instance of `base`. Returns slot info or None if unknown base. When a 2nd instance registers, slot 1 is renamed from 'base' to 'base-1' to prevent identity ambiguity. The rename info is returned as '_renamed_slot1'. + + `ready_gate=True` (opt-in): the instance starts in "starting" + instead of "active" - identity and token are issued (the CLI's MCP + config needs them before launch) but presence, mention routing, and + claim stay closed until the authenticated mark_ready transition. """ with self._lock: if base not in self._bases: @@ -237,8 +243,9 @@ def register(self, base: str, label: str | None = None) -> dict | None: # Fresh registrations are immediately authoritative. Identity # recovery/reclaim still uses chat_claim, but normal startup should - # not block on a manual confirmation step. - state = "active" + # not block on a manual confirmation step. Under the opt-in ready + # gate the instance instead enters "starting" (see register docstring). + state = "starting" if ready_gate else "active" inst = Instance(name=name, base=base, slot=slot, label=lbl, color=color, state=state) self._instances[name] = inst # Fresh registration (plus any slot-1 rename above) supersedes reclaimable @@ -270,36 +277,51 @@ def deregister(self, name: str, reclaimable: bool = False) -> dict | None: Returns result dict with 'ok' and optional '_renamed_back' info, or None if instance not found. """ + return self._remove_instance(name, reserve=True, reclaimable=reclaimable) + + def _remove_instance(self, name: str, *, reserve: bool = True, + reclaimable: bool = False, + require_state: str | None = None, + by_state: bool = False) -> dict | None: + """Shared removal path (deregister + cancel_starting + expire_crashed). + + The reservation is optional: the ready gate's cancel skips it so an + immediate relaunch reacquires the same name. Family rename-back + bookkeeping ALWAYS runs - a removal path that skipped it could strand a + 'base-1' name after its sibling died. `require_state` makes + check-and-remove atomic under one lock acquisition. + + `by_state` hands BOTH semantics to this method instead of the caller, + and adds a `cancelled` key to the result. Callers that do not use it + keep the documented `{'ok': True}` plus optional `_renamed_back` shape. + """ with self._lock: - if name not in self._instances: + inst_removed = self._instances.get(name) + if inst_removed is None: return None - inst_removed = self._instances[name] + if require_state is not None and inst_removed.state != require_state: + return None + cancelled = False + if by_state: + # Both branches are chosen here, holding the same lock that + # performs the removal. Reading the state in the caller and + # removing afterwards would let a wrapper re-gate in between, + # so an instance that is starting by the time it is removed + # would still take the reclaimable path. + cancelled = inst_removed.state == "starting" + reserve = not cancelled + reclaimable = not cancelled base = inst_removed.base del self._instances[name] - self._reserved[name] = time.time() + if reserve: + self._reserved[name] = time.time() if reclaimable: self._reclaimable[name] = inst_removed # token recoverable on reconnect (sleep) - # If family drops to 1 instance with a numbered name, rename back to base - renamed_back = None - family = [i for i in self._instances.values() if i.base == base] - if len(family) == 1: - remaining = family[0] - r_base, r_slot = self._parse_name(remaining.name) - if r_base == base and remaining.name != base: - old_name = remaining.name - del self._instances[old_name] - remaining.name = base - remaining.slot = 1 - base_cfg = self._bases.get(base, {}) - remaining.label = base_cfg.get("label", base.capitalize()) - remaining.color = _derive_color(base_cfg.get("color", "#888"), 1) - self._instances[base] = remaining - self._renames[old_name] = base - renamed_back = {"old": old_name, "new": base} + renamed_back = self._rename_back_if_single_locked(base) # A rename-back may have moved a live instance onto a (base, slot) that the - # just-deregistered identity still occupies in _reclaimable — drop such stale + # just-removed identity still occupies in _reclaimable — drop such stale # collisions so they can't outlive the live identity across a restart. self._evict_reclaimable_collisions_locked() @@ -307,10 +329,95 @@ def deregister(self, name: str, reclaimable: bool = False) -> dict | None: self._save_renames() self._save_instances() result = {"ok": True} + if by_state: + result["cancelled"] = cancelled if renamed_back: result["_renamed_back"] = renamed_back return result + def _rename_back_if_single_locked(self, base: str) -> dict | None: + """If the family drops to 1 instance with a numbered name, rename it + back to the bare base name. Caller holds self._lock.""" + renamed_back = None + family = [i for i in self._instances.values() if i.base == base] + if len(family) == 1: + remaining = family[0] + r_base, r_slot = self._parse_name(remaining.name) + if r_base == base and remaining.name != base: + old_name = remaining.name + del self._instances[old_name] + remaining.name = base + remaining.slot = 1 + base_cfg = self._bases.get(base, {}) + remaining.label = base_cfg.get("label", base.capitalize()) + remaining.color = _derive_color(base_cfg.get("color", "#888"), 1) + self._instances[base] = remaining + self._renames[old_name] = base + renamed_back = {"old": old_name, "new": base} + return renamed_back + + # --- Ready gate --- + + def get_state(self, name: str) -> str | None: + with self._lock: + inst = self._instances.get(name) + return inst.state if inst else None + + def mark_starting(self, name: str) -> bool: + """Re-enter the gate (CLI restart). Legal from starting|active only - + a pending (unclaimed) placeholder is not a gated launch.""" + with self._lock: + inst = self._instances.get(name) + if not inst or inst.state not in ("starting", "active"): + return False + inst.state = "starting" + self._notify() + self._save_instances() + return True + + def mark_ready(self, name: str) -> bool: + """The ONLY path that activates a starting instance.""" + with self._lock: + inst = self._instances.get(name) + if not inst or inst.state != "starting": + return False + inst.state = "active" + self._notify() + self._save_instances() + return True + + def cancel_starting(self, name: str) -> dict | None: + """Failed gate: remove WITHOUT the grace reservation and WITHOUT a + reclaimable entry, so an immediate relaunch reacquires the same name + and the dead token cannot be revived. + + Returns the removal record (falsy when nothing was removed) so callers + can act on `_renamed_back` exactly as they do for deregister; a cancelled + second instance renames its sibling back, and the caller still has to + migrate chat identity and history to the restored name. + """ + return self._remove_instance( + name, reserve=False, reclaimable=False, require_state="starting" + ) + + def expire_crashed(self, name: str) -> dict | None: + """Remove an instance whose wrapper has stopped heartbeating. + + The two cases need opposite treatment and the choice is state-dependent, + so it is made under the same lock that performs the removal. Reading the + state first and removing afterwards is a race: a wrapper re-gating from + active to starting in between would still be deregistered reclaimably, + which reserves its name and leaves a revivable token — exactly the ghost + the gate cancel exists to prevent. + + A `starting` instance is cancelled; anything else keeps the sleep/crash + behaviour so a machine waking up recovers its identity. Both semantics + are chosen inside the removal, not passed in. The returned record + carries `cancelled` so callers can report which path was taken; this is + the only removal that adds that key. + """ + return self._remove_instance(name, by_state=True) + # --- Identity Claim --- def claim(self, sender: str, target_name: str | None = None) -> dict | str: @@ -351,6 +458,11 @@ def claim(self, sender: str, target_name: str | None = None) -> dict | str: if not inst: error = f"No available {sender} instance. Is a wrapper registered?" + elif inst.state == "starting": + # Ready gate: only the authenticated mark_ready + # path may activate a starting instance - claim must not. + error = (f"{inst.name} is starting (ready gate); " + "claim is not allowed until it is ready") elif target_name is None or target_name == inst.name: # Accept current name — but don't auto-activate pending instances. # Pending instances must be named by human (lightbox) or reclaimed @@ -659,7 +771,12 @@ def resolve_token(self, token: str) -> dict | None: del self._reclaimable[name] changed = True continue - inst.state = "active" + # A recovered token must not skip the gate. An instance that was + # `starting` when the server stopped stays starting until an + # authenticated mark_ready; otherwise a restart would silently + # activate an agent whose CLI was never proven ready. + if inst.state != "starting": + inst.state = "active" self._instances[name] = inst del self._reclaimable[name] self._reserved.pop(name, None) diff --git a/tests/test_ready_gate_registry.py b/tests/test_ready_gate_registry.py new file mode 100644 index 00000000..e7683ca4 --- /dev/null +++ b/tests/test_ready_gate_registry.py @@ -0,0 +1,302 @@ +"""Ready-gate registry contract. + +The gate's registry invariants: gated registration enters `starting`; +`mark_ready` is the ONLY path that activates a starting instance; claim +refuses starting instances; cancel_starting removes with NO reservation and +NO reclaimable entry via the same removal primitive as deregister, so family +rename-back bookkeeping runs on both paths. +""" +import sys +import tempfile +import threading +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from registry import RuntimeRegistry + + +class ReadyGateRegistryTests(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.reg = RuntimeRegistry(data_dir=self._tmp.name) + self.reg.seed({"claude": {"label": "Claude", "color": "#da7756"}, + "codex": {"label": "Codex", "color": "#10a37f"}}) + + def tearDown(self): + self._tmp.cleanup() + + def test_gated_register_starts_in_starting(self): + result = self.reg.register("claude", ready_gate=True) + self.assertEqual(result["state"], "starting") + self.assertEqual(self.reg.get_state("claude"), "starting") + + def test_ungated_register_still_active(self): + result = self.reg.register("claude") + self.assertEqual(result["state"], "active") + self.assertEqual(self.reg.get_state("claude"), "active") + + def test_mark_ready_only_from_starting(self): + self.reg.register("claude", ready_gate=True) + self.assertTrue(self.reg.mark_ready("claude")) + self.assertEqual(self.reg.get_state("claude"), "active") + # Illegal transitions refuse + self.assertFalse(self.reg.mark_ready("claude")) # active -> ready: no + self.assertFalse(self.reg.mark_ready("gemini")) # unknown: no + self.reg.register("codex") # ungated, active + self.assertFalse(self.reg.mark_ready("codex")) # never entered gate: no + + def test_mark_starting_regates_active_but_not_pending(self): + self.reg.register("claude", ready_gate=True) + self.reg.mark_ready("claude") + self.assertTrue(self.reg.mark_starting("claude")) + self.assertEqual(self.reg.get_state("claude"), "starting") + # A pending instance (unclaimed placeholder) must NOT be re-gateable. + self.reg.register("codex") + with self.reg._lock: + self.reg._instances["codex"].state = "pending" + self.assertFalse(self.reg.mark_starting("codex")) + + def test_claim_refuses_starting_instance(self): + self.reg.register("claude", ready_gate=True) + res = self.reg.claim("claude") + self.assertIsInstance( + res, str, + "claim must return an error string, not activate a starting instance") + self.assertIn("starting", res) + self.assertEqual(self.reg.get_state("claude"), "starting") + + def test_cancel_starting_no_reservation_no_reclaim_dead_token(self): + first = self.reg.register("claude", ready_gate=True) + self.assertTrue(self.reg.cancel_starting("claude")) + self.assertIsNone(self.reg.resolve_token(first["token"])) + relaunch = self.reg.register("claude", ready_gate=True) + self.assertEqual(relaunch["name"], "claude") # not claude-2 + self.assertNotIn("_renamed_slot1", relaunch) + + def test_cancel_starting_second_instance_renames_back(self): + """Cancel shares deregister's family bookkeeping. claude active + + gated claude-2; cancelling claude-2 must rename claude-1 -> claude.""" + self.reg.register("claude") # slot 1, active + second = self.reg.register("claude", ready_gate=True) # slot1 -> claude-1 + self.assertEqual(second["name"], "claude-2") + self.assertTrue(self.reg.cancel_starting("claude-2")) + names = set(self.reg.get_all().keys()) + self.assertIn("claude", names, + "rename-back bookkeeping must run on cancel") + self.assertNotIn("claude-1", names) + + def test_cancel_starting_refuses_non_starting(self): + self.reg.register("claude") + self.assertFalse(self.reg.cancel_starting("claude")) # active: deregister only + + def test_normal_deregister_still_reserves(self): + """Pinned current behavior: the 30 s reservation SURVIVES normal deregister.""" + self.reg.register("claude") + self.reg.deregister("claude") + self.assertEqual(self.reg.register("claude")["name"], "claude-2") + + +class RemovalReturnShapeTests(unittest.TestCase): + """The gate must not change what the existing removals return. + + This branch is opt-in and promises unchanged behaviour when the gate is + unused. `cancelled` is meaningful only for the state-dependent expiry, so + the ordinary removals keep their documented `{'ok': True}` plus optional + `_renamed_back` exactly. + """ + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.reg = RuntimeRegistry(data_dir=self._tmp.name) + self.reg.seed({"claude": {"label": "Claude", "color": "#da7756"}}) + + def tearDown(self): + self._tmp.cleanup() + + def test_deregister_returns_exactly_ok(self): + self.reg.register("claude") + self.assertEqual(self.reg.deregister("claude"), {"ok": True}) + + def test_reclaimable_deregister_returns_exactly_ok(self): + self.reg.register("claude") + self.assertEqual(self.reg.deregister("claude", reclaimable=True), {"ok": True}) + + def test_cancel_starting_returns_exactly_ok(self): + self.reg.register("claude", ready_gate=True) + self.assertEqual(self.reg.cancel_starting("claude"), {"ok": True}) + + def test_deregister_with_rename_back_adds_only_that_key(self): + self.reg.register("claude") + self.reg.register("claude") # claude -> claude-1, claude-2 + result = self.reg.deregister("claude-2") + self.assertEqual(set(result), {"ok", "_renamed_back"}) + self.assertEqual(result["_renamed_back"], {"old": "claude-1", "new": "claude"}) + + def test_only_expire_crashed_reports_cancelled(self): + self.reg.register("claude", ready_gate=True) + self.assertIn("cancelled", self.reg.expire_crashed("claude")) + + +class ExpireCrashedAtomicityTests(unittest.TestCase): + """Expiry must decide by the state the instance has AT REMOVAL. + + The removal semantics differ by state, so choosing them outside the removal + lock is a race: a wrapper re-gating from active to starting in the window + would still be deregistered reclaimably, which reserves its name and leaves + a revivable token — the ghost the gate cancel exists to prevent. + """ + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.reg = RuntimeRegistry(data_dir=self._tmp.name) + self.reg.seed({"claude": {"label": "Claude", "color": "#da7756"}}) + + def tearDown(self): + self._tmp.cleanup() + + def test_starting_at_removal_time_is_cancelled(self): + inst = self.reg.register("claude") # registered active + self.assertTrue(self.reg.mark_starting("claude")) # wrapper re-gates + + record = self.reg.expire_crashed("claude") + + self.assertTrue(record["cancelled"]) + fresh = self.reg.register("claude", ready_gate=True) + self.assertEqual(fresh["name"], "claude", "must reacquire the bare name") + self.assertIsNone(self.reg.resolve_token(inst["token"]), + "the dead token must not be revivable") + + def test_active_at_removal_time_stays_reclaimable(self): + inst = self.reg.register("claude") + + record = self.reg.expire_crashed("claude") + + self.assertFalse(record["cancelled"]) + recovered = self.reg.resolve_token(inst["token"]) + self.assertIsNotNone(recovered, "a sleeping agent recovers its identity") + self.assertEqual(recovered["name"], "claude") + + def test_decision_is_not_taken_from_a_stale_read(self): + """The choice must not come from a state read before the removal. + + An implementation that reads the state first and removes afterwards acts + on a value that may already be wrong. The hook below fires only for such + an implementation — the atomic path never consults `get_state` — and it + moves the state inside exactly that window, so a split implementation + selects the wrong removal and fails here. + """ + self.reg.register("claude", ready_gate=True) # starting + real_get_state = self.reg.get_state + + def get_state_then_change(name): + observed = real_get_state(name) + self.reg.mark_ready(name) # the state moves in the window... + return observed # ...and the caller acts on the old one + + self.reg.get_state = get_state_then_change + + record = self.reg.expire_crashed("claude") + + self.assertIsNotNone(record, "expiry must remove the instance it was given") + self.assertTrue(record["cancelled"], + "a starting instance must take the cancel path") + + def test_a_concurrent_regate_cannot_land_inside_the_removal(self): + """Mutual exclusion, proven with a barrier inside the removal lock. + + The hook runs after expiry has chosen its semantics but before the lock + is released. A `mark_starting` issued from another thread must still be + waiting at that point; if it could complete, the decision and the + removal would be describing different states. + """ + self.reg.register("claude") + inside = threading.Event() + regate_done = threading.Event() + + original = self.reg._rename_back_if_single_locked + + def barrier(base): + inside.set() + regate_done.wait(0.5) + self.assertFalse( + regate_done.is_set(), + "mark_starting completed while the removal held the lock") + return original(base) + + self.reg._rename_back_if_single_locked = barrier + + def regate(): + inside.wait(1) + self.reg.mark_starting("claude") + regate_done.set() + + racer = threading.Thread(target=regate) + racer.start() + try: + self.assertIsNotNone(self.reg.expire_crashed("claude")) + finally: + racer.join(2) + self.assertFalse(racer.is_alive()) + + +class GateSurvivesRestartTests(unittest.TestCase): + """A server restart must not activate an agent whose CLI was never proven. + + On startup every persisted instance is reloaded as reclaimable, and a live + wrapper transparently recovers its identity through resolve_token. That + recovery path must preserve `starting`, or restarting the server silently + opens the gate for an agent that is still booting. + """ + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.seed = {"claude": {"label": "Claude", "color": "#da7756"}} + + def tearDown(self): + self._tmp.cleanup() + + def _restart(self): + """A fresh registry over the same data dir, as a server restart would.""" + reg = RuntimeRegistry(data_dir=self._tmp.name) + reg.seed(self.seed) + return reg + + def test_starting_survives_token_recovery_after_restart(self): + reg = RuntimeRegistry(data_dir=self._tmp.name) + reg.seed(self.seed) + inst = reg.register("claude", ready_gate=True) + + recovered = self._restart().resolve_token(inst["token"]) + + self.assertIsNotNone(recovered, "a live wrapper must still recover") + self.assertEqual(recovered["state"], "starting", + "restart must not activate an unproven CLI") + + def test_active_still_reactivates_after_restart(self): + """The sleep/crash recovery path for a proven agent is unchanged.""" + reg = RuntimeRegistry(data_dir=self._tmp.name) + reg.seed(self.seed) + inst = reg.register("claude") + + recovered = self._restart().resolve_token(inst["token"]) + + self.assertIsNotNone(recovered) + self.assertEqual(recovered["state"], "active") + + def test_ready_after_recovery_activates(self): + reg = RuntimeRegistry(data_dir=self._tmp.name) + reg.seed(self.seed) + inst = reg.register("claude", ready_gate=True) + + restarted = self._restart() + restarted.resolve_token(inst["token"]) + self.assertTrue(restarted.mark_ready("claude")) + self.assertEqual(restarted.get_state("claude"), "active") + + +if __name__ == "__main__": + unittest.main()