Skip to content

Commit dd900d2

Browse files
ampagentIsaac Miller
andcommitted
Preserve requested-route admission and unblock login cancellation
Retain the original route through credential renewal and re-resolve it at dispatch admission. Release cancelled login reservations on an independent executor even when login occupies the only default worker. Six behavioral regressions fail before and pass after; the full source suite and pinned conformance checks pass apart from the two approved orb connect-timeout exceptions. 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>
1 parent 72ced2f commit dd900d2

3 files changed

Lines changed: 62 additions & 22 deletions

File tree

‎lm15/login/manager.py‎

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
import secrets
3030
import threading
3131
import time
32-
from dataclasses import dataclass, field
32+
from dataclasses import dataclass, field, replace
3333
from datetime import datetime, timezone
3434
from typing import Any, Callable
3535

@@ -64,6 +64,7 @@ class ForgetResult:
6464
@dataclass(frozen=True, slots=True, kw_only=True)
6565
class _RequestSnapshot(RequestAuth):
6666
connection: Connection
67+
route: str
6768

6869

6970
@dataclass
@@ -778,6 +779,7 @@ def request_auth(self, provider: str, *, pinned: tuple[str, str] | None = None)
778779
client's ``(connection_id, generation)``; a mismatch is
779780
``connection_changed``, never a silent rebind (AUTH-20.1)."""
780781
provider = self.descriptor(provider).id
782+
route = provider
781783
document = self.store.read()
782784
slot, material = self._route_view(document, provider)
783785
provider = slot.provider
@@ -786,17 +788,17 @@ def request_auth(self, provider: str, *, pinned: tuple[str, str] | None = None)
786788
# have died mid-exchange. Only the lock can tell: wait for it,
787789
# re-read, and reuse the sibling's result; a marker still there
788790
# once we hold the lock is an interrupted renewal (AUTH-20.4).
789-
return self._renew(provider, pinned)
791+
return replace(self._renew(provider, pinned), route=route)
790792
self._check_selected(provider, slot, material, pinned)
791793
flow = flow_for_material(provider, material) # type: ignore[arg-type]
792794
assert material is not None
793795
expiry = flow.expiry(material)
794796
if expiry in ("never", None):
795-
return self._auth_from(provider, flow, material, slot)
797+
return replace(self._auth_from(provider, flow, material, slot), route=route)
796798
now_ms = int(self._clock() * 1000)
797799
if now_ms < expiry - self._lead_ms(flow.lifetime_s(material)):
798-
return self._auth_from(provider, flow, material, slot)
799-
return self._renew(provider, pinned)
800+
return replace(self._auth_from(provider, flow, material, slot), route=route)
801+
return replace(self._renew(provider, pinned), route=route)
800802

801803
def _check_selected(self, provider: str, slot: _Slot, material: Material | None, pinned: tuple[str, str] | None) -> None:
802804
if pinned is not None and (slot.connection_id != pinned[0] or str(slot.generation) != pinned[1]):
@@ -832,14 +834,14 @@ def _check_selected(self, provider: str, slot: _Slot, material: Material | None,
832834
provider=provider, connection_id=slot.connection_id,
833835
)
834836

835-
def _auth_from(self, provider: str, flow: Any, material: Material, slot: _Slot) -> RequestAuth:
837+
def _auth_from(self, provider: str, flow: Any, material: Material, slot: _Slot) -> _RequestSnapshot:
836838
try:
837839
value = flow.request_auth(material, slot.settings)
838840
connection = slot.connection()
839841
assert connection is not None
840842
return _RequestSnapshot(
841843
credential=value.credential, headers=dict(value.headers), base_url=value.base_url,
842-
account_id=value.account_id, named=value.named, connection=connection,
844+
account_id=value.account_id, named=value.named, connection=connection, route=provider,
843845
)
844846
except LoginDenied as exc:
845847
raise AuthOperationError(f"{provider}: {exc}", reason="login_required", stage="resolution",
@@ -854,11 +856,11 @@ def _admit(self, snapshot: _RequestSnapshot) -> None:
854856
"""
855857
connection = snapshot.connection
856858
with self.store.transaction() as txn:
857-
slot, material = self._view(txn.read(), connection.provider)
858-
self._check_selected(connection.provider, slot, material,
859+
slot, material = self._route_view(txn.read(), snapshot.route)
860+
self._check_selected(snapshot.route, slot, material,
859861
(connection.id, connection.identity_generation))
860862

861-
def _renew(self, provider: str, pinned: tuple[str, str] | None) -> RequestAuth:
863+
def _renew(self, provider: str, pinned: tuple[str, str] | None) -> _RequestSnapshot:
862864
"""AUTH-20.4: lock, re-read, reuse a sibling's fresh result, else
863865
mark in-flight, exchange, write — all under the lock."""
864866
ctx = LoginContext(ui=_NoUI(), deadline=self._monotonic() + 60.0, provider=provider, clock=self._monotonic,
@@ -1024,22 +1026,28 @@ async def status(self, provider: str):
10241026

10251027
async def login(self, provider: str, method: Any = None, **kwargs: Any):
10261028
import asyncio
1029+
from concurrent.futures import ThreadPoolExecutor
10271030

10281031
cancel = kwargs.get("cancel") or threading.Event()
10291032
kwargs["cancel"] = cancel
10301033
try:
10311034
return await self._run(self.sync.login, provider, method, **kwargs)
10321035
except asyncio.CancelledError:
10331036
cancel.set()
1034-
cleanup = asyncio.create_task(self._run(self.sync._cancel_login_attempt, provider, cancel))
1035-
# Repeated task cancellation must not let a reservation outlive
1036-
# CancelledError. Do not wait for the provider's worker to return.
1037-
while not cleanup.done():
1038-
try:
1039-
await asyncio.shield(cleanup)
1040-
except asyncio.CancelledError:
1041-
continue
1042-
cleanup.result()
1037+
# A blocking login can occupy every default-executor worker.
1038+
# Durable cancellation must not queue behind those workers.
1039+
with ThreadPoolExecutor(max_workers=1) as executor:
1040+
cleanup = asyncio.get_running_loop().run_in_executor(
1041+
executor, self.sync._cancel_login_attempt, provider, cancel,
1042+
)
1043+
# Repeated cancellation must not let the reservation outlive
1044+
# CancelledError. The provider's worker need not return first.
1045+
while not cleanup.done():
1046+
try:
1047+
await asyncio.shield(cleanup)
1048+
except asyncio.CancelledError:
1049+
continue
1050+
cleanup.result()
10431051
raise
10441052

10451053
async def set_api_key(self, provider: str, key: str, **kwargs: Any):

‎tests/test_login.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -461,28 +461,35 @@ def test_one_active_attempt_per_slot(sandbox: Path) -> None:
461461
@pytest.mark.parametrize("cooperative", [False, True])
462462
def test_async_cancellation_cannot_replace_saved_connection(sandbox: Path, monkeypatch, cooperative) -> None:
463463
import asyncio
464+
from concurrent.futures import ThreadPoolExecutor
464465
from lm15.login import AsyncAuth
465466
from lm15.login.flows.base import LoginResult
466467
from lm15.login.flows.xai import XaiFlow
467468

468-
entered, approval = threading.Event(), threading.Event()
469+
entered, approval = asyncio.Event(), threading.Event()
469470
auth = make_auth(sandbox, Clock())
470471
original = auth.set_api_key("xai", "original")
471472

472473
def delayed_login(self, ctx, method, settings, answers):
473-
entered.set()
474+
loop.call_soon_threadsafe(entered.set)
474475
assert approval.wait(5)
475476
if cooperative:
476477
ctx.check()
477478
return LoginResult(material={"type": "api_key", "key": "replacement"}, label="replacement")
478479

479480
monkeypatch.setattr(XaiFlow, "login", delayed_login)
481+
loop = None
480482

481483
async def run():
484+
nonlocal loop
485+
loop = asyncio.get_running_loop()
486+
loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
482487
task = asyncio.create_task(AsyncAuth(auth).login("xai", "device", ui=ScriptUI(), replace=original.id))
483488
try:
484-
assert await asyncio.to_thread(entered.wait, 3)
489+
await asyncio.wait_for(entered.wait(), 3)
485490
task.cancel()
491+
done, _ = await asyncio.wait({task}, timeout=1)
492+
assert task in done, "cancellation must finish before the blocked login worker returns"
486493
with pytest.raises(asyncio.CancelledError):
487494
await task
488495
reopened = Auth.local(sandbox / "credentials.json")

‎tests/test_managed_dispatch.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,31 @@ def after_encode(self, **kwargs):
148148
assert transport.requests == []
149149

150150

151+
@pytest.mark.parametrize("asynchronous", [False, True])
152+
@pytest.mark.parametrize("change", ["replace", "logout"])
153+
def test_shared_route_change_before_handoff_sends_nothing(monkeypatch, asynchronous, change):
154+
auth = Auth.memory()
155+
auth.set_api_key("meta", "shared")
156+
transport = transport_for("complete", asynchronous)
157+
lm = (AsyncLMRouter if asynchronous else LMRouter)(RouterConfig(auth=auth, transport=transport)).lm("meta-chat:muse")
158+
emit = BaseProviderLM._emit
159+
160+
def after_encode(self, **kwargs):
161+
sent = emit(self, **kwargs)
162+
assert dict(sent.headers)["Authorization"] == "Bearer shared"
163+
auth.set_api_key("meta-chat", "separate")
164+
if change == "logout":
165+
auth.logout("meta-chat")
166+
return sent
167+
168+
monkeypatch.setattr(BaseProviderLM, "_emit", after_encode)
169+
with pytest.raises(AuthOperationError) as error:
170+
invoke(lm, "complete", asynchronous, model="muse")
171+
assert error.value.reason == ("connection_changed" if change == "replace" else "login_required")
172+
assert auth.status("meta").connection is not None
173+
assert transport.requests == []
174+
175+
151176
def test_concurrent_logout_of_encoded_request(monkeypatch):
152177
auth = Auth.memory()
153178
local(auth, "a")

0 commit comments

Comments
 (0)