diff --git a/backend/apps/nine_router/lent_credential_refresh.py b/backend/apps/nine_router/lent_credential_refresh.py new file mode 100644 index 00000000..74e219c6 --- /dev/null +++ b/backend/apps/nine_router/lent_credential_refresh.py @@ -0,0 +1,93 @@ +"""Keeping this device usable while the cloud holds custody of a provider credential. + +Handing a credential to the cloud strips the local refresh token on purpose: exactly one +holder may rotate it. 9Router's refresh dispatcher then bails on the falsy refreshToken +without calling the provider, so nothing renews the access token and every local call +starts failing a few hours after the handover, with no error that explains why. + +This is the other half of that trade. The cloud rotates; this device asks the cloud for a +fresh access token shortly before the current one dies. Only lent connections are touched: +one that still holds a refresh token is 9Router's job and must be left alone. + +Deliberately lazy. Each pull rewrites db.json, which means stopping the router and starting +it again, so we act only inside the margin and never on a fixed cadence. +""" +from __future__ import annotations + +import asyncio +import logging +import time +from typing import List, Optional + +from typeguard import typechecked + +from backend.apps.nine_router import credential_lease, credential_store + +logger = logging.getLogger(__name__) + +# Pull this far ahead of expiry. Wide enough that a failure leaves room for several retries +# before anything 401s, narrow enough that we are not restarting the router for fun. +REFRESH_MARGIN_S = 15 * 60 +CHECK_INTERVAL_S = 120.0 +# A dead network would otherwise mean a router restart every two minutes forever. +FAILURE_BACKOFF_S = 900.0 + + +@typechecked +def p_seconds_left(expires_at: Optional[str]) -> Optional[float]: + """Seconds until this token dies, or None when the timestamp is unreadable.""" + ms = credential_lease.expires_ms(expires_at) + if ms <= 0: + return None + return (ms / 1000.0) - time.time() + + +@typechecked +def lent_connections_needing_a_pull() -> List[str]: + """Connections the cloud holds whose access token is spent or nearly so. + + A missing refresh token is what marks a connection as lent, and an unreadable expiry is + treated as due: better one wasted pull than a token that quietly stops working. + """ + due: List[str] = [] + for connection_id in credential_store.list_oauth_connection_ids(): + cred = credential_store.read_credential(connection_id) + if cred is None or cred.refresh_token: + continue + left = p_seconds_left(cred.expires_at) + if left is None or left <= REFRESH_MARGIN_S: + due.append(connection_id) + return due + + +@typechecked +async def refresh_lent_credentials() -> int: + """Top up every lent connection that needs it. Returns how many are now good.""" + refreshed = 0 + for connection_id in lent_connections_needing_a_pull(): + outcome = await credential_lease.pull_access_token(connection_id) + if outcome.status == "refreshed": + refreshed += 1 + continue + # Never the token itself, only why we could not get one. + logger.warning( + "could not renew the cloud-held credential %s: %s %s", + connection_id, + outcome.status, + outcome.detail, + ) + return refreshed + + +@typechecked +async def lent_credential_loop() -> None: + while True: + delay = CHECK_INTERVAL_S + try: + due = lent_connections_needing_a_pull() + if due and await refresh_lent_credentials() == 0: + delay = FAILURE_BACKOFF_S + except Exception: + logger.exception("lent-credential refresh pass failed") + delay = FAILURE_BACKOFF_S + await asyncio.sleep(delay) diff --git a/backend/apps/workflows/cloud/routes.py b/backend/apps/workflows/cloud/routes.py index 2d664671..49743afb 100644 --- a/backend/apps/workflows/cloud/routes.py +++ b/backend/apps/workflows/cloud/routes.py @@ -17,6 +17,7 @@ from pydantic import BaseModel, ConfigDict from typeguard import typechecked from backend.apps.workflows import storage +from backend.apps.nine_router.lent_credential_refresh import lent_credential_loop from backend.apps.workflows.cloud import client as cloud from backend.apps.workflows.cloud.handover import TargetOutcome, hand_to_cloud, take_back from backend.apps.workflows.cloud.run_files import LocalRunFile, described, downloads_root, fetch_missing @@ -27,7 +28,14 @@ from backend.config.Apps import SubApp @asynccontextmanager async def cloud_workflows_lifespan(): - yield + # Lending a credential upward strips this device's ability to renew it, so something has to + # ask the cloud for a fresh one before the old one dies. Without this, turning on a cloud + # workflow quietly stops local agents a few hours later. + task = asyncio.create_task(lent_credential_loop()) + try: + yield + finally: + task.cancel() cloud_workflows = SubApp("cloud_workflows", cloud_workflows_lifespan) diff --git a/backend/tests/test_lent_credential_refresh.py b/backend/tests/test_lent_credential_refresh.py new file mode 100644 index 00000000..bc5f35e9 --- /dev/null +++ b/backend/tests/test_lent_credential_refresh.py @@ -0,0 +1,234 @@ +"""Local work must survive the cloud borrowing your login. + +Lending strips this device's refresh token so only one holder can rotate. 9Router's refresh +dispatcher then bails on the falsy refreshToken WITHOUT calling the provider, so nothing +renews the access token: a few hours after a successful handover every local agent starts +failing, with no error that names the cause. This file pins the other half of the trade. + +The two ways to get it wrong are opposite and both bad: never pulling (local dies), and +pulling constantly (every pull stops and restarts the router, so a 2-minute cadence would +make the app unusable). +""" +import time + +import pytest + +from backend.apps.nine_router import lent_credential_refresh as lcr +from backend.apps.nine_router.credential_lease import LeaseOutcome +from backend.apps.nine_router.credential_store import ProviderCredential + + +def p_iso(seconds_from_now: float) -> str: + from datetime import datetime, timezone + + return datetime.fromtimestamp(time.time() + seconds_from_now, tz=timezone.utc).isoformat() + + +@pytest.fixture +def p_connections(monkeypatch): + """Install a fake 9router db. `refresh` present means this device still owns it.""" + + def install(rows): + creds = { + r["id"]: ProviderCredential( + connection_id=r["id"], + provider=r.get("provider", "claude"), + access_token="at", + refresh_token=r.get("refresh"), + expires_at=r.get("expires"), + ) + for r in rows + } + monkeypatch.setattr(lcr.credential_store, "list_oauth_connection_ids", lambda: list(creds)) + monkeypatch.setattr(lcr.credential_store, "read_credential", lambda cid: creds.get(cid)) + return creds + + return install + + +@pytest.fixture +def p_pull(monkeypatch): + def install(*statuses): + seq = list(statuses) + calls = [] + + async def fake(connection_id: str) -> LeaseOutcome: + calls.append(connection_id) + return LeaseOutcome(status=seq.pop(0) if seq else "refreshed") + + monkeypatch.setattr(lcr.credential_lease, "pull_access_token", fake) + return calls + + return install + + +def test_a_connection_this_device_still_owns_is_never_touched(p_connections): + # It has its own refresh token, so 9Router renews it. Pulling would fight the router for no reason. + p_connections([{"id": "mine", "refresh": "rt", "expires": p_iso(30)}]) + assert lcr.lent_connections_needing_a_pull() == [] + + +def test_a_lent_connection_with_hours_left_is_left_alone(p_connections): + # Every pull costs a router stop and start, so acting early would be worse than acting late. + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(6 * 3600)}]) + assert lcr.lent_connections_needing_a_pull() == [] + + +def test_a_lent_connection_inside_the_margin_is_due(p_connections): + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(lcr.REFRESH_MARGIN_S - 60)}]) + assert lcr.lent_connections_needing_a_pull() == ["lent"] + + +def test_an_already_expired_lent_connection_is_due(p_connections): + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(-3600)}]) + assert lcr.lent_connections_needing_a_pull() == ["lent"] + + +def test_an_unreadable_expiry_is_treated_as_due(p_connections): + # One wasted pull beats a token that silently stops working because we could not read a date. + p_connections([{"id": "lent", "refresh": None, "expires": "not-a-date"}]) + assert lcr.lent_connections_needing_a_pull() == ["lent"] + + +def test_a_missing_expiry_is_treated_as_due(p_connections): + p_connections([{"id": "lent", "refresh": None, "expires": None}]) + assert lcr.lent_connections_needing_a_pull() == ["lent"] + + +def test_only_the_due_lent_ones_are_selected_out_of_a_mixed_set(p_connections): + p_connections([ + {"id": "mine", "refresh": "rt", "expires": p_iso(10)}, + {"id": "lent-fresh", "refresh": None, "expires": p_iso(4 * 3600)}, + {"id": "lent-due", "refresh": None, "expires": p_iso(60)}, + {"id": "lent-dead", "refresh": None, "expires": p_iso(-99)}, + ]) + assert lcr.lent_connections_needing_a_pull() == ["lent-due", "lent-dead"] + + +@pytest.mark.asyncio +async def test_a_due_connection_gets_pulled(p_connections, p_pull): + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(-1)}]) + calls = p_pull("refreshed") + assert await lcr.refresh_lent_credentials() == 1 + assert calls == ["lent"] + + +@pytest.mark.asyncio +async def test_nothing_due_means_no_router_restart(p_connections, p_pull): + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(6 * 3600)}]) + calls = p_pull("refreshed") + assert await lcr.refresh_lent_credentials() == 0 + assert calls == [], "a pull rewrites db.json and bounces the router; do not do it for nothing" + + +@pytest.mark.asyncio +async def test_an_offline_pull_fails_without_raising(p_connections, p_pull): + # Being offline is normal. It must not take the loop down, and it must not be silent either. + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(-1)}]) + p_pull("cloud_rejected") + assert await lcr.refresh_lent_credentials() == 0 + + +@pytest.mark.asyncio +async def test_one_dead_connection_does_not_block_the_others(p_connections, p_pull): + p_connections([ + {"id": "a", "refresh": None, "expires": p_iso(-1)}, + {"id": "b", "refresh": None, "expires": p_iso(-1)}, + ]) + calls = p_pull("cloud_rejected", "refreshed") + assert await lcr.refresh_lent_credentials() == 1 + assert calls == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_signed_out_user_is_reported_not_retried_into_a_storm(p_connections, p_pull, caplog): + p_connections([{"id": "lent", "refresh": None, "expires": p_iso(-1)}]) + p_pull("not_signed_in") + assert await lcr.refresh_lent_credentials() == 0 + + +@pytest.mark.asyncio +async def test_a_failure_never_writes_a_token_into_the_log(monkeypatch): + # Not caplog: backend/main.py pins propagate=False on the 'backend' logger and caplog listens + # at the root, so these records only exist if you sit on the logger itself. + import io + import logging + + # A distinctive secret, so this cannot pass by luck. + secret = "sk-ant-oat01-NEVER-LOG-ME-9f3c2b" + monkeypatch.setattr( + lcr.credential_store, + "list_oauth_connection_ids", + lambda: ["lent"], + ) + monkeypatch.setattr( + lcr.credential_store, + "read_credential", + lambda cid: ProviderCredential( + connection_id="lent", provider="claude", access_token=secret, + refresh_token=None, expires_at=p_iso(-1), + ), + ) + + async def leaky(connection_id: str) -> LeaseOutcome: + return LeaseOutcome(status="cloud_rejected", detail="HTTP 500") + + monkeypatch.setattr(lcr.credential_lease, "pull_access_token", leaky) + buf = io.StringIO() + handler = logging.StreamHandler(buf) + handler.setLevel(logging.WARNING) + lcr.logger.addHandler(handler) + try: + await lcr.refresh_lent_credentials() + finally: + lcr.logger.removeHandler(handler) + + written = buf.getvalue() + assert secret not in written, "the access token must never reach a log line" + assert "lent" in written, "but which connection failed has to be diagnosable" + + +# The join. Everything above passes even if nothing ever runs the loop, which is exactly how +# lease_to_cloud sat unwired for its whole life. + +@pytest.mark.asyncio +async def test_the_cloud_subsystem_actually_starts_the_refresh_loop(monkeypatch): + import asyncio + + from backend.apps.workflows.cloud import routes + + started = asyncio.Event() + + async def fake_loop(): + started.set() + await asyncio.sleep(3600) + + monkeypatch.setattr(routes, "lent_credential_loop", fake_loop) + + async with routes.cloud_workflows_lifespan(): + await asyncio.wait_for(started.wait(), timeout=2.0) + + +@pytest.mark.asyncio +async def test_shutting_the_subsystem_down_stops_the_loop(monkeypatch): + import asyncio + + from backend.apps.workflows.cloud import routes + + running = asyncio.Event() + cancelled = asyncio.Event() + + async def fake_loop(): + running.set() + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + cancelled.set() + raise + + monkeypatch.setattr(routes, "lent_credential_loop", fake_loop) + + async with routes.cloud_workflows_lifespan(): + await asyncio.wait_for(running.wait(), timeout=2.0) + await asyncio.sleep(0) + assert cancelled.is_set(), "a leaked task would keep bouncing the router after shutdown"