Keep managed routes cancellation and verification tied to their connection

Signed-off-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-01a0c095-4917-7617-9dac-e99f1ef22d41
Co-authored-by: Isaac Miller <isaac@cmpnd.ai>

Author: ampagent

lm15/login/manager.py

@@ -158,7 +158,7 @@ def __init__(self, store: Store, *, clock: Callable[[], float] = time.time,
         self._monotonic = monotonic
         self._opener = opener
         self._sleep = sleep
-        self._active: dict[str, threading.Event] = {}  # provider -> cancel event of a login in this process
+        self._active: dict[str, tuple[str, threading.Event]] = {}  # provider -> (attempt id, cancel event)
         self._lock = threading.RLock()
         self._closed = False
 
@@ -213,6 +213,20 @@ def _view(self, document: dict[str, Any], provider: str) -> tuple[_Slot, Materia
             )
         return slot, material
 
+    def _route_view(self, document: dict[str, Any], provider: str) -> tuple[_Slot, Material | None]:
+        """An explicit route slot wins; otherwise use its advertised owner."""
+        slots = (document.get(META_KEY) or {}).get("slots") or {}
+        if provider in document or provider in slots:
+            return self._view(document, provider)
+        owners = [owner for owner, record in slots.items()
+                  if provider in record.get("routes", ()) and provider in self.descriptor(owner).routes]
+        if len(owners) > 1:
+            raise AuthOperationError(
+                f"{provider}: multiple saved connections advertise this route; select a connection",
+                reason="connection_changed", stage="resolution", recovery="select_connection", provider=provider,
+            )
+        return self._view(document, owners[0] if owners else provider)
+
     @staticmethod
     def _put(document: dict[str, Any], slot: _Slot, material: Material | None) -> dict[str, Any]:
         meta = document.setdefault(META_KEY, {"version": STORE_VERSION, "slots": {}})
@@ -248,7 +262,7 @@ def connections(self) -> tuple[Connection, ...]:
 
     def status(self, provider: str) -> ConnectionStatus:
         provider = self.descriptor(provider).id
-        slot, material = self._view(self.store.read(), provider)
+        slot, material = self._route_view(self.store.read(), provider)
         connection = slot.connection()
         verification = None
         if slot.verification:
@@ -332,10 +346,8 @@ def login(
         # Reservation (AUTH-17/18): storage proven writable, one active
         # attempt per slot, generation observed — before any browser opens.
         self.store.reserve()
-        expected = self._reserve(provider, attempt_id, replace, lifetime_s)
-        with self._lock:
-            self._active[provider] = cancel
         try:
+            expected = self._reserve(provider, attempt_id, replace, lifetime_s, cancel=cancel)
             for field_ in chosen.fields:
                 if field_.id not in answers:
                     if not field_.required and field_.type != "select":
@@ -379,6 +391,7 @@ def login(
             except LM15Error:
                 self._release(provider, attempt_id)
                 raise
+            ctx.check()
             return self._commit(provider, attempt_id, expected, chosen, result, settings)
         except BaseException:
             # Any exit without a saved connection ends the attempt — Ctrl-C
@@ -390,7 +403,9 @@ def login(
             raise
         finally:
             with self._lock:
-                self._active.pop(provider, None)
+                active = self._active.get(provider)
+                if active is not None and active[0] == attempt_id:
+                    self._active.pop(provider)
 
     def _choose_method(self, descriptor: ProviderDescriptor, method: str | LoginMethod | None, ui: AuthUI, *,
                        allow_unverified: bool) -> LoginMethod:
@@ -440,7 +455,8 @@ def _choose_method(self, descriptor: ProviderDescriptor, method: str | LoginMeth
             reason="invalid_login_state", stage="interaction", recovery="choose_method", provider=descriptor.id,
         )
 
-    def _reserve(self, provider: str, attempt_id: str, replace: str | None, lifetime_s: float) -> int:
+    def _reserve(self, provider: str, attempt_id: str, replace: str | None, lifetime_s: float,
+                 *, cancel: threading.Event | None = None) -> int:
         """Reserve the slot's single active attempt; return the generation
         the commit must find."""
         now = self._clock()
@@ -473,6 +489,11 @@ def reserve(document: dict[str, Any]) -> dict[str, Any] | None:
                 )
             slot.attempt = {"id": attempt_id, "expected_generation": str(slot.generation),
                             "started_at_s": now, "lifetime_s": float(lifetime_s)}
+            if cancel is not None:
+                with self._lock:
+                    if cancel.is_set():
+                        raise LoginCancelled("login cancelled before reservation")
+                    self._active[provider] = (attempt_id, cancel)
             return self._put(document, slot, material)
 
         document = self.store.mutate(reserve)
@@ -552,17 +573,33 @@ def cancel(document: dict[str, Any]) -> dict[str, Any] | None:
             if not slot.attempt:
                 outcome["result"] = "complete" if slot.connection_id else "none"
                 return None
+            outcome["attempt_id"] = slot.attempt["id"]
             slot.attempt = None
             outcome["result"] = "cancelled"
             return self._put(document, slot, material)
 
         self.store.mutate(cancel)  # the durable record first (AUTH-19) …
         with self._lock:
-            event = self._active.get(provider)
-        if event is not None:
-            event.set()  # … then the running attempt in this process
+            active = self._active.get(provider)
+        if active is not None and active[0] == outcome.get("attempt_id"):
+            active[1].set()  # … then the running attempt in this process
         return outcome["result"]
 
+    def _cancel_login_attempt(self, provider: str, event: threading.Event) -> None:
+        """Cancel only this async caller's reservation, ordered against commit."""
+        provider = self.descriptor(provider).id
+
+        def cancel(document: dict[str, Any]) -> dict[str, Any] | None:
+            slot, material = self._view(document, provider)
+            with self._lock:
+                active = self._active.get(provider)
+            if active is None or active[1] is not event or not slot.attempt or slot.attempt.get("id") != active[0]:
+                return None
+            slot.attempt = None
+            return self._put(document, slot, material)
+
+        self.store.mutate(cancel)
+
     # ─── setup without a provider round-trip (AUTH-17) ─────────────────
 
     def set_api_key(self, provider: str, key: str, *, replace: str | None = None) -> Connection:
@@ -650,9 +687,9 @@ def forget(document: dict[str, Any]) -> dict[str, Any] | None:
             )
             outcome.update(forgot=True, generation=new.generation, routes=new.routes)
             with self._lock:
-                event = self._active.get(provider)
-            if event is not None:
-                event.set()
+                active = self._active.get(provider)
+            if active is not None:
+                active[1].set()
             return self._put(document, new, material_after)
 
         self.store.mutate(forget)
@@ -673,7 +710,9 @@ def _resolve_target(self, provider_or_connection: str) -> tuple[str, str | None]
                 "no saved connection has that id", reason="attempt_unavailable", stage="resolution",
                 recovery="select_connection",
             )
-        return self.descriptor(provider_or_connection).id, None
+        provider = self.descriptor(provider_or_connection).id
+        slot, _ = self._route_view(self.store.read(), provider)
+        return slot.provider, None
 
     # ─── verification (AUTH-17) ────────────────────────────────────────
 
@@ -689,11 +728,21 @@ def verify(self, provider: str, *, router_config: Any = None) -> Verification:
         if definition is None or not definition.supports.models:
             return Verification(result="unverified", check="models", detail="this route has no safe non-inference check")
         from ..router import LMRouter, RouterConfig
+        from .bound import pinned_auth
+        from .types import ModelSelection
 
         config = router_config or RouterConfig()
         from dataclasses import replace as _replace
 
-        router = LMRouter(_replace(config, auth=self))
+        slot, _ = self._route_view(self.store.read(), provider)
+        connection = slot.connection()
+        managed = self
+        if connection is not None:
+            selection = ModelSelection(provider=provider, model="verify", connection_id=connection.id,
+                                       identity_generation=connection.identity_generation)
+            managed = pinned_auth(self, selection)
+        # Verify the saved connection, never an explicit router-config key.
+        router = LMRouter(_replace(config, auth=managed, api_keys={}, credentials={}))
         checked = _now_iso(self._clock)
         try:
             lm = router.lm(f"{provider}:verify")
@@ -703,13 +752,14 @@ def verify(self, provider: str, *, router_config: Any = None) -> Verification:
             result = Verification(result="rejected", checked_at=checked, check="models", detail=exc.code)
         finally:
             router.close()
-        self._record_verification(provider, result)
+        if connection is not None:
+            self._record_verification(connection, result)
         return result
 
-    def _record_verification(self, provider: str, result: Verification) -> None:
+    def _record_verification(self, connection: Connection, result: Verification) -> None:
         def record(document: dict[str, Any]) -> dict[str, Any] | None:
-            slot, material = self._view(document, provider)
-            if slot.connection_id is None:
+            slot, material = self._view(document, connection.provider)
+            if slot.connection_id != connection.id or str(slot.generation) != connection.identity_generation:
                 return None
             slot.verification = {"result": result.result, "checked_at": result.checked_at, "check": result.check,
                                  "detail": result.detail}
@@ -729,7 +779,8 @@ def request_auth(self, provider: str, *, pinned: tuple[str, str] | None = None)
         ``connection_changed``, never a silent rebind (AUTH-20.1)."""
         provider = self.descriptor(provider).id
         document = self.store.read()
-        slot, material = self._view(document, provider)
+        slot, material = self._route_view(document, provider)
+        provider = slot.provider
         if slot.renewal_in_flight and slot.state != "indeterminate":
             # A sibling may be renewing right now (it holds the lock), or may
             # have died mid-exchange. Only the lock can tell: wait for it,
@@ -901,7 +952,7 @@ def close(self) -> None:
         with self._lock:
             self._closed = True
             events = list(self._active.values())
-        for event in events:
+        for _, event in events:
             event.set()
 
     def _check_open(self) -> None:
@@ -972,7 +1023,24 @@ async def status(self, provider: str):
         return await self._run(self.sync.status, provider)
 
     async def login(self, provider: str, method: Any = None, **kwargs: Any):
-        return await self._run(self.sync.login, provider, method, **kwargs)
+        import asyncio
+
+        cancel = kwargs.get("cancel") or threading.Event()
+        kwargs["cancel"] = cancel
+        try:
+            return await self._run(self.sync.login, provider, method, **kwargs)
+        except asyncio.CancelledError:
+            cancel.set()
+            cleanup = asyncio.create_task(self._run(self.sync._cancel_login_attempt, provider, cancel))
+            # Repeated task cancellation must not let a reservation outlive
+            # CancelledError. Do not wait for the provider's worker to return.
+            while not cleanup.done():
+                try:
+                    await asyncio.shield(cleanup)
+                except asyncio.CancelledError:
+                    continue
+            cleanup.result()
+            raise
 
     async def set_api_key(self, provider: str, key: str, **kwargs: Any):
         return await self._run(self.sync.set_api_key, provider, key, **kwargs)

tests/test_login.py

@@ -458,6 +458,48 @@ def test_one_active_attempt_per_slot(sandbox: Path) -> None:
     auth.login("xai", "device", ui=ScriptUI())
 
 
+@pytest.mark.parametrize("cooperative", [False, True])
+def test_async_cancellation_cannot_replace_saved_connection(sandbox: Path, monkeypatch, cooperative) -> None:
+    import asyncio
+    from lm15.login import AsyncAuth
+    from lm15.login.flows.base import LoginResult
+    from lm15.login.flows.xai import XaiFlow
+
+    entered, approval = threading.Event(), threading.Event()
+    auth = make_auth(sandbox, Clock())
+    original = auth.set_api_key("xai", "original")
+
+    def delayed_login(self, ctx, method, settings, answers):
+        entered.set()
+        assert approval.wait(5)
+        if cooperative:
+            ctx.check()
+        return LoginResult(material={"type": "api_key", "key": "replacement"}, label="replacement")
+
+    monkeypatch.setattr(XaiFlow, "login", delayed_login)
+
+    async def run():
+        task = asyncio.create_task(AsyncAuth(auth).login("xai", "device", ui=ScriptUI(), replace=original.id))
+        try:
+            assert await asyncio.to_thread(entered.wait, 3)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            reopened = Auth.local(sandbox / "credentials.json")
+            assert reopened.status("xai").connection.id == original.id
+            # Cancellation releases the reservation before approval returns,
+            # so a new operation can replace A immediately.
+            replacement = reopened.set_api_key("xai", "newer", replace=original.id)
+            return replacement.id
+        finally:
+            approval.set()
+
+    replacement_id = asyncio.run(run())  # joins the cancelled task's worker
+    reopened = Auth.local(sandbox / "credentials.json")
+    assert reopened.status("xai").connection.id == replacement_id
+    assert reopened.store.read()["xai"]["key"] == "newer"
+
+
 def test_an_interrupt_ends_the_attempt_and_frees_the_slot(sandbox: Path) -> None:
     # Ctrl-C — a notebook's cancel — while waiting for the provider, or at
     # a prompt: the attempt is over, so its reservation must not block the
@@ -969,3 +1011,57 @@ def test_meta_without_a_plan_says_pay_as_you_go(sandbox: Path) -> None:
                                                  "is_subs_active": False}))
     connection = auth.login("meta", "device", ui=ScriptUI(), allow_unverified=True)
     assert connection.account_label == "no Muse Code subscription (pay-as-you-go)"
+
+
+@pytest.mark.parametrize("route", ["meta", "meta-chat", "meta-anthropic"])
+def test_meta_routes_share_connection_renewal_and_logout(sandbox: Path, route) -> None:
+    from lm15.login.flows.meta import KEY_LIFETIME_S
+
+    clock = Clock()
+    server = FakeMeta({"api_key": "key-a", "base_url": "https://eu.api.meta.ai/v1"})
+    auth = make_auth(sandbox, clock, server)
+    connection = auth.login("meta", "device", ui=ScriptUI(), allow_unverified=True)
+    transport = FakeTransport([_models_reply(), _models_reply()])
+    router = LMRouter(RouterConfig(auth=auth, transport=transport, env={"META_API_KEY": "payg"}))
+    lm = router.lm(f"{route}:muse-wonder-full-1.0")
+    assert lm.list_models()
+    server.mint = {"api_key": "key-b", "base_url": "https://us.api.meta.ai/v1"}
+    clock.advance(KEY_LIFETIME_S)
+    assert lm.list_models()
+    for sent, region, key in zip(transport.requests, ("eu", "us"), ("key-a", "key-b")):
+        assert sent.url.split("?", 1)[0] == f"https://{region}.api.meta.ai/v1/models"
+        headers = {k.lower(): v for k, v in sent.headers}
+        assert key in (headers.get("x-api-key"), headers.get("authorization", "").removeprefix("Bearer "))
+    assert auth.status(route).connection.id == connection.id
+    assert auth.status(route).connection.credential_revision == "2"
+    assert auth.logout(route).forgot
+    with pytest.raises(AuthOperationError, match="signed out"):
+        lm.list_models()
+    assert len(transport.requests) == 2
+    router.close()
+
+
+@pytest.mark.parametrize("replace_connection", [False, True])
+@pytest.mark.parametrize("rejected", [False, True])
+def test_verification_stays_with_checked_connection(sandbox: Path, replace_connection, rejected) -> None:
+    auth = make_auth(sandbox, Clock())
+    original = auth.set_api_key("openai-chat", "key-a")
+    reply = FakeResponse(401, b'{"error":{"message":"rejected"}}') if rejected else _models_reply()
+
+    class ReplacingTransport(FakeTransport):
+        def stream(self, request):
+            if replace_connection:
+                auth.set_api_key("openai-chat", "key-b", replace=original.id)
+            return super().stream(request)
+
+    transport = ReplacingTransport([reply])
+    result = auth.verify("openai-chat", router_config=RouterConfig(transport=transport))
+    assert result.result == ("rejected" if rejected else "valid")
+    assert dict(transport.requests[0].headers)["Authorization"] == "Bearer key-a"
+    status = Auth.local(sandbox / "credentials.json").status("openai-chat")
+    if replace_connection:
+        assert status.connection.id != original.id
+        assert status.verification is None
+    else:
+        assert status.connection.id == original.id
+        assert status.verification == result