Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

hermes_cli/nous_auth_keepalive.py

@@ -0,0 +1,310 @@
+"""Background keepalive for long-lived Nous Portal sessions."""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from typing import Optional
+
+from hermes_cli.auth import (
+    ACCESS_TOKEN_REFRESH_SKEW_SECONDS,
+    NOUS_INVOKE_JWT_MIN_TTL_SECONDS,
+    AuthError,
+    _agent_key_is_usable,
+    _is_expiring,
+    get_provider_auth_state,
+    resolve_nous_runtime_credentials,
+)
+
+logger = logging.getLogger(__name__)
+
+# Two things have to line up for the keepalive to actually keep anything alive.
+#
+# 1. The tick has to be frequent enough to see the credential before it dies.
+#    Nous credential lifetimes are not fixed: this varies by account, and
+#    installs have been observed at both ~3594s and ~899s. The tick therefore
+#    derives from the lifetime the server actually issued rather than assuming
+#    an hour, capped by the configured interval and floored so a pathological
+#    lifetime cannot spin the thread.
+#
+# 2. The refresh has to fire while the tick can still act on it. The refresh
+#    only triggers once a credential is within a skew window of expiry, so a
+#    tick spaced wider than that window steps straight over it. The keepalive
+#    widens the window to "will this credential outlive my next tick?" instead
+#    of the request-path default of 120s. Without this, ticking faster only
+#    narrows the gap; it never closes it.
+#
+# The original 6-hour tick against a one-hour credential failed both tests, so
+# in practice every hour expired reactively into a 401 plus a re-auth retry.
+NOUS_AUTH_KEEPALIVE_INTERVAL_SECONDS = 15 * 60
+NOUS_AUTH_KEEPALIVE_MIN_INTERVAL_SECONDS = 60
+# Ticks per credential lifetime. Four keeps the refresh comfortably ahead of
+# expiry without making the thread chatty.
+NOUS_AUTH_KEEPALIVE_TICKS_PER_LIFETIME = 4
+NOUS_AUTH_KEEPALIVE_INITIAL_DELAY_SECONDS = 60
+NOUS_AUTH_KEEPALIVE_INTERVAL_CONFIG_KEY = "keepalive_interval_seconds"
+
+_keepalive_lock = threading.Lock()
+_keepalive_stop = threading.Event()
+_keepalive_thread: Optional[threading.Thread] = None
+
+
+def _timeout_seconds(value: Optional[float]) -> float:
+    if value is not None:
+        return float(value)
+    try:
+        return float(os.getenv("HERMES_NOUS_TIMEOUT_SECONDS", "15"))
+    except (TypeError, ValueError):
+        return 15.0
+
+
+def _nous_config() -> dict:
+    """Return the ``nous:`` section of config.yaml, or {} on any failure.
+
+    Imported lazily: this module is loaded by the gateway and the web server
+    during startup, and the config loader pulls in a wider dependency graph
+    than the keepalive itself needs.
+    """
+    try:
+        from hermes_cli.config import load_config
+
+        section = load_config().get("nous")
+        return section if isinstance(section, dict) else {}
+    except Exception:
+        return {}
+
+
+def _interval_seconds(value: Optional[int]) -> int:
+    """Resolve the keepalive tick interval.
+
+    Explicit argument wins, then ``nous.keepalive_interval_seconds`` in
+    config.yaml, then the module default. This is a behavioural threshold
+    rather than a credential, so it lives in config.yaml and not in .env.
+    A non-positive result disables the keepalive thread entirely, which is
+    the documented way to turn it off.
+    """
+    if value is not None:
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            return NOUS_AUTH_KEEPALIVE_INTERVAL_SECONDS
+
+    raw = _nous_config().get(NOUS_AUTH_KEEPALIVE_INTERVAL_CONFIG_KEY)
+    if raw is None or (isinstance(raw, str) and not raw.strip()):
+        return NOUS_AUTH_KEEPALIVE_INTERVAL_SECONDS
+    try:
+        return int(float(raw))
+    except (TypeError, ValueError):
+        logger.warning(
+            "Ignoring invalid nous.%s=%r; using %ds",
+            NOUS_AUTH_KEEPALIVE_INTERVAL_CONFIG_KEY,
+            raw,
+            NOUS_AUTH_KEEPALIVE_INTERVAL_SECONDS,
+        )
+        return NOUS_AUTH_KEEPALIVE_INTERVAL_SECONDS
+
+
+def _observed_lifetime_seconds() -> Optional[int]:
+    """Lifetime the server issued for the current Nous credentials, in seconds.
+
+    Both the access token and the invoke agent key carry their own lifetime and
+    they are not always equal, so the shorter one governs. Returns None when no
+    usable value is stored, in which case the caller keeps its configured tick.
+    """
+    state = get_provider_auth_state("nous") or {}
+    lifetimes = []
+    for key in ("expires_in", "agent_key_expires_in"):
+        try:
+            value = int(float(state.get(key)))
+        except (TypeError, ValueError):
+            continue
+        if value > 0:
+            lifetimes.append(value)
+    return min(lifetimes) if lifetimes else None
+
+
+def _tick_seconds(configured_interval: int, lifetime: Optional[int]) -> int:
+    """Tick fast enough to refresh several times per credential lifetime."""
+    if not lifetime or lifetime <= 0:
+        return configured_interval
+    derived = lifetime // NOUS_AUTH_KEEPALIVE_TICKS_PER_LIFETIME
+    return max(
+        NOUS_AUTH_KEEPALIVE_MIN_INTERVAL_SECONDS,
+        min(configured_interval, derived),
+    )
+
+
+def _refresh_horizon_seconds(tick_seconds: int, floor_seconds: int) -> int:
+    """How much life a credential needs to be left alone this tick.
+
+    A credential that will not survive until the next tick has to be refreshed
+    now, because nothing will look at it again before it expires. Hence
+    tick + skew rather than the request path's bare skew.
+    """
+    return max(floor_seconds, tick_seconds + ACCESS_TOKEN_REFRESH_SKEW_SECONDS)
+
+
+def _entry_state(entry: object) -> dict:
+    return {
+        "agent_key": getattr(entry, "agent_key", None),
+        "agent_key_expires_at": getattr(entry, "agent_key_expires_at", None),
+        "scope": getattr(entry, "scope", None),
+    }
+
+
+def _refresh_selected_pool_entry(
+    *,
+    min_key_ttl_seconds: int,
+    min_access_ttl_seconds: Optional[int] = None,
+) -> Optional[bool]:
+    """Refresh the current Nous credential pool entry when it is stale.
+
+    Returns True when a pool entry exists and is usable/refreshed, False when a
+    pool exists but no entry can be used, and None when no Nous pool exists.
+    """
+    try:
+        from agent.credential_pool import load_pool
+
+        pool = load_pool("nous")
+    except Exception as exc:
+        logger.debug("Nous auth keepalive: credential pool unavailable: %s", exc)
+        return None
+
+    if not pool or not pool.has_credentials():
+        return None
+
+    try:
+        entry = pool.select()
+    except Exception as exc:
+        logger.debug("Nous auth keepalive: credential pool selection failed: %s", exc)
+        return False
+
+    if entry is None:
+        return False
+
+    if min_access_ttl_seconds is None:
+        min_access_ttl_seconds = ACCESS_TOKEN_REFRESH_SKEW_SECONDS
+    access_expiring = _is_expiring(
+        getattr(entry, "expires_at", None),
+        min_access_ttl_seconds,
+    )
+    key_usable = _agent_key_is_usable(_entry_state(entry), min_key_ttl_seconds)
+    if access_expiring or not key_usable:
+        refreshed = pool.try_refresh_current()
+        if refreshed is None:
+            return False
+        logger.debug("Nous auth keepalive: refreshed credential pool entry")
+        return True
+
+    return True
+
+
+def refresh_nous_auth_keepalive_once(
+    *,
+    min_key_ttl_seconds: int = NOUS_INVOKE_JWT_MIN_TTL_SECONDS,
+    min_access_ttl_seconds: Optional[int] = None,
+    timeout_seconds: Optional[float] = None,
+) -> bool:
+    """Refresh Nous auth once if credentials are configured."""
+    min_key_ttl_seconds = max(60, int(min_key_ttl_seconds))
+
+    pool_result = _refresh_selected_pool_entry(
+        min_key_ttl_seconds=min_key_ttl_seconds,
+        min_access_ttl_seconds=min_access_ttl_seconds,
+    )
+    if pool_result is not None:
+        return pool_result
+
+    state = get_provider_auth_state("nous")
+    if not state:
+        return False
+
+    try:
+        resolve_nous_runtime_credentials(
+            timeout_seconds=_timeout_seconds(timeout_seconds),
+        )
+        logger.debug("Nous auth keepalive: refreshed singleton auth state")
+        return True
+    except AuthError as exc:
+        if exc.relogin_required:
+            logger.info("Nous auth keepalive requires re-login: %s", exc)
+        else:
+            logger.debug("Nous auth keepalive failed: %s", exc)
+        return False
+    except Exception as exc:
+        logger.debug("Nous auth keepalive failed: %s", exc)
+        return False
+
+
+def _keepalive_loop(
+    stop_event: threading.Event,
+    *,
+    interval_seconds: int,
+    initial_delay_seconds: int,
+    min_key_ttl_seconds: int,
+    timeout_seconds: Optional[float],
+) -> None:
+    if initial_delay_seconds > 0 and stop_event.wait(initial_delay_seconds):
+        return
+
+    while not stop_event.is_set():
+        # Re-read each pass: the lifetime can change when the account, plan, or
+        # server-side policy does, and a keepalive that caches it would go stale
+        # in exactly the case it exists to cover.
+        tick = _tick_seconds(interval_seconds, _observed_lifetime_seconds())
+        horizon = _refresh_horizon_seconds(tick, min_key_ttl_seconds)
+        refresh_nous_auth_keepalive_once(
+            min_key_ttl_seconds=horizon,
+            min_access_ttl_seconds=horizon,
+            timeout_seconds=timeout_seconds,
+        )
+        stop_event.wait(tick)
+
+
+def start_nous_auth_keepalive(
+    *,
+    interval_seconds: Optional[int] = None,
+    initial_delay_seconds: int = NOUS_AUTH_KEEPALIVE_INITIAL_DELAY_SECONDS,
+    min_key_ttl_seconds: int = NOUS_INVOKE_JWT_MIN_TTL_SECONDS,
+    timeout_seconds: Optional[float] = None,
+) -> Optional[threading.Thread]:
+    """Start the process-wide Nous auth keepalive thread."""
+    interval_seconds = _interval_seconds(interval_seconds)
+    if interval_seconds <= 0:
+        return None
+
+    global _keepalive_thread
+    with _keepalive_lock:
+        if _keepalive_thread is not None and _keepalive_thread.is_alive():
+            return _keepalive_thread
+
+        _keepalive_stop.clear()
+        _keepalive_thread = threading.Thread(
+            target=_keepalive_loop,
+            args=(_keepalive_stop,),
+            kwargs={
+                "interval_seconds": int(interval_seconds),
+                "initial_delay_seconds": max(0, int(initial_delay_seconds)),
+                "min_key_ttl_seconds": max(60, int(min_key_ttl_seconds)),
+                "timeout_seconds": timeout_seconds,
+            },
+            daemon=True,
+            name="nous-auth-keepalive",
+        )
+        _keepalive_thread.start()
+        logger.debug("Nous auth keepalive started")
+        return _keepalive_thread
+
+
+def stop_nous_auth_keepalive(timeout: float = 5.0) -> None:
+    """Stop the keepalive thread. Intended for graceful shutdown/tests."""
+    global _keepalive_thread
+    with _keepalive_lock:
+        thread = _keepalive_thread
+        _keepalive_stop.set()
+    if thread is not None and thread.is_alive():
+        thread.join(timeout=timeout)
+    with _keepalive_lock:
+        if _keepalive_thread is thread:
+            _keepalive_thread = None