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()