522 lines
21 KiB
Python
522 lines
21 KiB
Python
"""OAuth credential storage and refresh for the Honcho memory provider.
|
|
|
|
An access token authenticates like a scoped API key, so it is stored as the
|
|
host's ``apiKey``; the refresh token is exchanged before expiry to keep it live.
|
|
|
|
Refresh tokens rotate with single-use reuse detection: replaying a stale token
|
|
revokes the grant, so every refresh persists the rotated token atomically and is
|
|
serialized (in-process lock + cross-process file lock). A failed exchange never
|
|
raises into the agent: transient failures retry once immediately (the replay
|
|
grace window is short); a permanent error such as invalid_grant marks the grant
|
|
dead so nothing keeps hitting the endpoint and callers surface a re-login
|
|
prompt. A server-side 401 on a locally-valid token is recovered via
|
|
``force_refresh_token``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import threading
|
|
import time
|
|
from contextlib import contextmanager, suppress
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ACCESS_TOKEN_PREFIX = "hch-at-"
|
|
REFRESH_TOKEN_PREFIX = "hch-rt-"
|
|
|
|
_REFRESH_SKEW_SECONDS = 120 # refresh this early so an in-flight request never races expiry
|
|
_REFRESH_TIMEOUT_SECONDS = 15.0 # short: sits on the path to a memory call
|
|
_REFRESH_RETRY_DELAY_SECONDS = 2.0 # replayed refresh tokens are honored only briefly after rotation
|
|
# Cap for one exchange cycle (attempt + pause + retry): it runs under the global
|
|
# refresh locks, so it must not hold them for two full HTTP timeouts.
|
|
_REFRESH_TOTAL_BUDGET_SECONDS = 20.0
|
|
# After a transient failure, fail open without re-exchanging for this long so N
|
|
# waiting threads don't serialize N exchange cycles against a failing endpoint.
|
|
_REFRESH_FAILURE_COOLDOWN_SECONDS = 30.0
|
|
# OAuth error codes a retry can never fix — the grant itself is dead.
|
|
_PERMANENT_OAUTH_ERRORS = frozenset({"invalid_grant", "invalid_client", "unauthorized_client"})
|
|
|
|
# Derived from the canonical prefixes so a prefix change can't silently break redaction.
|
|
_TOKEN_VALUE_RE = re.compile(
|
|
rf"({re.escape(ACCESS_TOKEN_PREFIX)}|{re.escape(REFRESH_TOKEN_PREFIX)})[A-Za-z0-9._~+/=-]+"
|
|
)
|
|
|
|
|
|
def redact_tokens(text: str) -> str:
|
|
"""Replace any embedded token values with their prefix plus a placeholder."""
|
|
return _TOKEN_VALUE_RE.sub(lambda m: f"{m.group(1)}[redacted]", text)
|
|
|
|
|
|
_redact_tokens = redact_tokens # backward-compat alias for older importers
|
|
|
|
|
|
class OAuthRefreshError(Exception):
|
|
"""Token endpoint rejected the refresh. ``permanent`` means re-login is required."""
|
|
|
|
def __init__(self, message: str, *, error: str = "", permanent: bool = False):
|
|
super().__init__(message)
|
|
self.error = error
|
|
self.permanent = permanent
|
|
|
|
|
|
# Serializes refresh across threads; state is re-checked under it so racing
|
|
# callers don't replay a rotated refresh token.
|
|
_refresh_lock = threading.Lock()
|
|
|
|
|
|
def _os_lock(fh, lock: bool) -> None:
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
fh.seek(0)
|
|
msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK if lock else msvcrt.LK_UNLCK, 1)
|
|
else:
|
|
import fcntl
|
|
|
|
fcntl.flock(fh.fileno(), fcntl.LOCK_EX if lock else fcntl.LOCK_UN)
|
|
|
|
|
|
@contextmanager
|
|
def _config_refresh_lock(path: Path):
|
|
"""Machine-wide advisory lock (``<config>.lock``) around read-refresh-persist:
|
|
a sibling process sharing this honcho.json must not replay the single-use
|
|
refresh token. Best-effort — without flock it degrades to in-process only."""
|
|
fh = None
|
|
try:
|
|
lock_path = Path(f"{path}.lock")
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
fh = open(lock_path, "a+b")
|
|
_os_lock(fh, True)
|
|
except Exception:
|
|
logger.debug("Honcho OAuth cross-process lock unavailable; in-process only", exc_info=True)
|
|
if fh is not None:
|
|
fh.close()
|
|
fh = None
|
|
try:
|
|
yield
|
|
finally:
|
|
if fh is not None:
|
|
with suppress(Exception):
|
|
_os_lock(fh, False)
|
|
fh.close()
|
|
|
|
|
|
# Per-grant state keyed by (config path, host); single-key dict ops are atomic
|
|
# under the GIL, so no separate lock.
|
|
# (expires_at, access): lets the hot path skip the honcho.json read while the
|
|
# token is well clear of expiry. A stale entry can't break auth (the token stays
|
|
# valid until its own expiry); it only defers noticing out-of-band rotation.
|
|
_expiry_cache: dict[tuple[str, str], tuple[float, str]] = {}
|
|
# sha256 of the permanently rejected refresh token; a re-login rotates the
|
|
# token, so the digest comparison self-clears.
|
|
_dead_grants: dict[tuple[str, str], str] = {}
|
|
# monotonic time of the last transient exchange failure (drives the fail-open cooldown).
|
|
_refresh_failure_at: dict[tuple[str, str], float] = {}
|
|
# (config mtime_ns, verdict): reauth_required only changes when the file is
|
|
# rewritten, so an unchanged mtime skips the parse.
|
|
_reauth_check_cache: dict[tuple[str, str], tuple[int, bool]] = {}
|
|
|
|
|
|
def _in_failure_cooldown(key: tuple[str, str]) -> bool:
|
|
failed_at = _refresh_failure_at.get(key)
|
|
return failed_at is not None and (time.monotonic() - failed_at) < _REFRESH_FAILURE_COOLDOWN_SECONDS
|
|
|
|
|
|
def _refresh_token_digest(cred: OAuthCredential) -> str:
|
|
return hashlib.sha256(cred.refresh_token.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _grant_is_dead(key: tuple[str, str], cred: OAuthCredential) -> bool:
|
|
return _dead_grants.get(key) == _refresh_token_digest(cred)
|
|
|
|
|
|
def _mark_grant_dead(key: tuple[str, str], cred: OAuthCredential) -> None:
|
|
_dead_grants[key] = _refresh_token_digest(cred)
|
|
_reauth_check_cache.pop(key, None) # verdict changed without a config rewrite
|
|
|
|
|
|
def _load_cred(path: Path, host: str, raw: dict[str, Any] | None = None) -> OAuthCredential | None:
|
|
"""Credential from ``host``'s block in ``raw`` (or the file at ``path``)."""
|
|
source = raw if raw is not None else _read_config(path)
|
|
return OAuthCredential.from_host_block((source.get("hosts") or {}).get(host) or {})
|
|
|
|
|
|
def reauth_required(path: Path, host: str) -> bool:
|
|
"""True when ``host``'s OAuth grant is dead and only a new login fixes it."""
|
|
key = (str(path), host)
|
|
if key not in _dead_grants:
|
|
return False
|
|
try:
|
|
mtime = path.stat().st_mtime_ns
|
|
except OSError:
|
|
mtime = -1
|
|
cached = _reauth_check_cache.get(key)
|
|
if cached is not None and cached[0] == mtime:
|
|
return cached[1]
|
|
cred = _load_cred(path, host)
|
|
result = cred is not None and _grant_is_dead(key, cred)
|
|
_reauth_check_cache[key] = (mtime, result)
|
|
return result
|
|
|
|
|
|
def any_dead_grants() -> bool:
|
|
"""Cheap predicate letting hot-path callers skip config resolution when healthy."""
|
|
return bool(_dead_grants)
|
|
|
|
|
|
def _num(value, cast):
|
|
"""``cast(value)``, or ``cast(0)`` when the stored value is malformed."""
|
|
try:
|
|
return cast(value)
|
|
except (TypeError, ValueError):
|
|
return cast(0)
|
|
|
|
|
|
def is_oauth_access_token(value: str | None) -> bool:
|
|
"""True when ``value`` is an OAuth access token (vs a static API key)."""
|
|
return bool(value) and value.startswith(ACCESS_TOKEN_PREFIX)
|
|
|
|
|
|
@dataclass
|
|
class OAuthCredential:
|
|
"""An OAuth grant as stored in a honcho.json host block: ``access_token`` is
|
|
the host's ``apiKey``, the rest lives in its ``oauth`` sub-block.
|
|
``expires_at`` is absolute epoch seconds."""
|
|
|
|
access_token: str
|
|
refresh_token: str
|
|
expires_at: float
|
|
client_id: str
|
|
token_endpoint: str
|
|
scope: str = "write"
|
|
token_type: str = "Bearer"
|
|
# Transient consent peer name — set only on a fresh grant, never persisted.
|
|
consent_peer_name: str | None = None
|
|
|
|
@classmethod
|
|
def from_host_block(cls, block: dict[str, Any]) -> "OAuthCredential | None":
|
|
"""Build a credential from a honcho.json host block, or None if incomplete."""
|
|
oauth = block.get("oauth")
|
|
access = block.get("apiKey")
|
|
if not isinstance(oauth, dict) or not is_oauth_access_token(access):
|
|
return None
|
|
refresh, endpoint, client_id = oauth.get("refreshToken"), oauth.get("tokenEndpoint"), oauth.get("clientId")
|
|
if not (refresh and endpoint and client_id):
|
|
return None
|
|
return cls(
|
|
access_token=access,
|
|
refresh_token=str(refresh),
|
|
expires_at=_num(oauth.get("expiresAt", 0), float),
|
|
client_id=str(client_id),
|
|
token_endpoint=str(endpoint),
|
|
scope=str(oauth.get("scope", "write")),
|
|
token_type=str(oauth.get("tokenType", "Bearer")),
|
|
)
|
|
|
|
@classmethod
|
|
def from_token_response(
|
|
cls,
|
|
body: dict[str, Any],
|
|
*,
|
|
now: float,
|
|
client_id: str,
|
|
token_endpoint: str,
|
|
scope: str = "write",
|
|
token_type: str = "Bearer",
|
|
what: str = "grant",
|
|
) -> "OAuthCredential":
|
|
"""Build a credential from an OAuth token response; ``expires_in`` is relative to ``now``."""
|
|
access, refresh = body.get("access_token"), body.get("refresh_token")
|
|
if not is_oauth_access_token(access) or not refresh:
|
|
raise ValueError(f"{what} missing access_token/refresh_token")
|
|
return cls(
|
|
access_token=access,
|
|
refresh_token=str(refresh),
|
|
expires_at=now + _num(body.get("expires_in", 0), int),
|
|
client_id=client_id,
|
|
token_endpoint=token_endpoint,
|
|
scope=str(body.get("scope", scope)),
|
|
token_type=str(body.get("token_type", token_type)),
|
|
)
|
|
|
|
def oauth_block(self) -> dict[str, Any]:
|
|
"""The ``oauth`` sub-block to persist (the access token lives in apiKey)."""
|
|
return {
|
|
"refreshToken": self.refresh_token,
|
|
"expiresAt": int(self.expires_at),
|
|
"clientId": self.client_id,
|
|
"tokenEndpoint": self.token_endpoint,
|
|
"scope": self.scope,
|
|
"tokenType": self.token_type,
|
|
}
|
|
|
|
def is_expired(self, *, now: float, skew: float = _REFRESH_SKEW_SECONDS) -> bool:
|
|
"""True when the access token is within ``skew`` seconds of expiry."""
|
|
return now >= (self.expires_at - skew)
|
|
|
|
|
|
# HTTP indirection: tests monkeypatch these module attributes, and the refresh
|
|
# path looks them up at call time.
|
|
def _http_json(method: str, url: str, *, data=None, timeout: float, strict: bool = True) -> tuple[int, Any]:
|
|
"""Return ``(status, parsed JSON body)``. ``strict`` raises on non-2xx /
|
|
non-JSON; otherwise a 4xx passes through (RFC 8628 polling reads the OAuth
|
|
error off a 400) and a non-JSON / non-object body parses to ``{}``."""
|
|
import httpx
|
|
|
|
resp = httpx.request(method, url, data=data, timeout=timeout)
|
|
if strict:
|
|
resp.raise_for_status()
|
|
return resp.status_code, resp.json()
|
|
try:
|
|
body = resp.json()
|
|
except ValueError:
|
|
body = {}
|
|
return resp.status_code, body if isinstance(body, dict) else {}
|
|
|
|
|
|
def _http_post_form_status(url: str, data: dict[str, str], timeout: float) -> tuple[int, dict[str, Any]]:
|
|
"""POST form-encoded ``data``; return ``(status, body)`` without raising on 4xx."""
|
|
return _http_json("POST", url, data=data, timeout=timeout, strict=False)
|
|
|
|
|
|
def _exchange_refresh_token(
|
|
cred: OAuthCredential, *, now: float, timeout: float = _REFRESH_TIMEOUT_SECONDS
|
|
) -> OAuthCredential:
|
|
"""Run the refresh_token grant and return the rotated credential. Raises
|
|
``OAuthRefreshError`` (with the endpoint's error body) on an error response,
|
|
transport errors as-is; callers fail open."""
|
|
status, body = _http_post_form_status(
|
|
cred.token_endpoint,
|
|
{"grant_type": "refresh_token", "client_id": cred.client_id, "refresh_token": cred.refresh_token},
|
|
timeout,
|
|
)
|
|
if status >= 400:
|
|
error = str(body.get("error") or "")
|
|
description = str(body.get("error_description") or "")
|
|
detail = " — ".join(p for p in (error, description) if p) or "no error body"
|
|
raise OAuthRefreshError(
|
|
_redact_tokens(f"token endpoint returned HTTP {status}: {detail}"),
|
|
error=error,
|
|
permanent=error in _PERMANENT_OAUTH_ERRORS,
|
|
)
|
|
return OAuthCredential.from_token_response(
|
|
body, now=now, client_id=cred.client_id, token_endpoint=cred.token_endpoint,
|
|
scope=cred.scope, token_type=cred.token_type, what="refresh response",
|
|
)
|
|
|
|
|
|
def _exchange_with_retry(cred: OAuthCredential, *, now: float) -> OAuthCredential:
|
|
"""Exchange the refresh token, retrying once on transient failure. The retry
|
|
cannot wait (replay grace window is short) and the cycle is capped by
|
|
``_REFRESH_TOTAL_BUDGET_SECONDS`` because it runs under the refresh locks."""
|
|
deadline = time.monotonic() + _REFRESH_TOTAL_BUDGET_SECONDS
|
|
try:
|
|
return _exchange_refresh_token(cred, now=now)
|
|
except OAuthRefreshError as exc:
|
|
if exc.permanent:
|
|
raise
|
|
first: Exception = exc
|
|
except Exception as exc:
|
|
first = exc
|
|
remaining = deadline - time.monotonic() - _REFRESH_RETRY_DELAY_SECONDS
|
|
if remaining <= 0:
|
|
raise first
|
|
logger.warning("Honcho OAuth token exchange failed, retrying once: %s", _redact_tokens(str(first)))
|
|
time.sleep(_REFRESH_RETRY_DELAY_SECONDS)
|
|
return _exchange_refresh_token(cred, now=now, timeout=min(remaining, _REFRESH_TIMEOUT_SECONDS))
|
|
|
|
|
|
def _rotate_and_persist(
|
|
path: Path,
|
|
host: str,
|
|
key: tuple[str, str],
|
|
cred: OAuthCredential,
|
|
*,
|
|
now: float,
|
|
op_label: str = "refresh",
|
|
) -> OAuthCredential | None:
|
|
"""Exchange ``cred`` and persist the rotation; ``None`` on failure (logged).
|
|
A permanent OAuth error marks the grant dead so later calls skip the endpoint
|
|
until a new login rotates the refresh token."""
|
|
try:
|
|
rotated = _exchange_with_retry(cred, now=now)
|
|
except OAuthRefreshError as exc:
|
|
if exc.permanent:
|
|
_mark_grant_dead(key, cred)
|
|
logger.error(
|
|
"Honcho OAuth grant for host %s is no longer valid (%s); "
|
|
"run 'hermes honcho setup' to re-authenticate", host, exc,
|
|
)
|
|
return None
|
|
_refresh_failure_at[key] = time.monotonic()
|
|
logger.warning("Honcho OAuth %s failed for host %s: %s", op_label, host, exc)
|
|
return None
|
|
except Exception as exc:
|
|
_refresh_failure_at[key] = time.monotonic()
|
|
logger.warning("Honcho OAuth %s failed for host %s: %s", op_label, host, _redact_tokens(str(exc)))
|
|
return None
|
|
_persist_credential(path, host, rotated)
|
|
return rotated
|
|
|
|
|
|
def _read_config(path: Path) -> dict[str, Any]:
|
|
try:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return {}
|
|
|
|
|
|
def _atomic_write_config(path: Path, raw: dict[str, Any]) -> None:
|
|
"""Write ``raw`` to ``path`` atomically with 0600 on the new file."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_name(f".{path.name}.tmp")
|
|
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
|
fh.write(json.dumps(raw, indent=2) + "\n")
|
|
except Exception:
|
|
tmp.unlink(missing_ok=True)
|
|
raise
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def _deep_merge(base: dict[str, Any], overlay: dict[str, Any]) -> dict[str, Any]:
|
|
"""Recursively merge ``overlay`` into ``base`` (overlay wins on scalars/lists)."""
|
|
for key, value in overlay.items():
|
|
if isinstance(value, dict) and isinstance(base.get(key), dict):
|
|
_deep_merge(base[key], value)
|
|
else:
|
|
base[key] = value
|
|
return base
|
|
|
|
|
|
def _persist_credential(path: Path, host: str, cred: OAuthCredential, raw: dict[str, Any] | None = None) -> None:
|
|
"""Write ``cred`` into ``host``'s block (apiKey + oauth) of ``raw`` (default:
|
|
the file's current content), leaving the rest intact; marks the grant live."""
|
|
raw = _read_config(path) if raw is None else raw
|
|
block = raw.setdefault("hosts", {}).setdefault(host, {})
|
|
block["apiKey"] = cred.access_token
|
|
block["oauth"] = cred.oauth_block()
|
|
_atomic_write_config(path, raw)
|
|
key = (str(path), host)
|
|
_expiry_cache[key] = (cred.expires_at, cred.access_token)
|
|
_dead_grants.pop(key, None)
|
|
_refresh_failure_at.pop(key, None)
|
|
|
|
|
|
def ensure_fresh_token(
|
|
path: Path,
|
|
host: str,
|
|
raw: dict[str, Any] | None = None,
|
|
*,
|
|
now: float | None = None,
|
|
) -> tuple[str | None, bool]:
|
|
"""Return ``(access_token, refreshed)`` for ``host``, refreshing if near expiry.
|
|
``(None, False)`` when the host has no OAuth credential (plain API key).
|
|
Refresh failures are swallowed: the current (possibly stale) token comes back
|
|
with ``refreshed=False``; a permanently rejected grant is marked dead and the
|
|
session's 401 recovery escalates it to the user."""
|
|
now = time.time() if now is None else now
|
|
key = (str(path), host)
|
|
|
|
# Hot path: trust the cached expiry while well clear of the skew window (no
|
|
# disk read). Bypassed when an explicit ``raw`` is supplied.
|
|
if raw is None:
|
|
cached = _expiry_cache.get(key)
|
|
if cached is not None and now < cached[0] - _REFRESH_SKEW_SECONDS:
|
|
return cached[1], False
|
|
|
|
cred = _load_cred(path, host, raw)
|
|
if cred is None:
|
|
_expiry_cache.pop(key, None)
|
|
return None, False
|
|
|
|
_expiry_cache[key] = (cred.expires_at, cred.access_token)
|
|
if not cred.is_expired(now=now) or _in_failure_cooldown(key):
|
|
return cred.access_token, False
|
|
|
|
with _refresh_lock, _config_refresh_lock(path):
|
|
# Re-read under both locks: another thread or process may have just
|
|
# rotated the token — adopt theirs instead of replaying the old one.
|
|
current = _load_cred(path, host) or cred
|
|
if not current.is_expired(now=now):
|
|
return current.access_token, current.access_token != cred.access_token
|
|
# The lock holder we waited on may have just failed; fail open too.
|
|
if _grant_is_dead(key, current) or _in_failure_cooldown(key):
|
|
return current.access_token, False
|
|
rotated = _rotate_and_persist(path, host, key, current, now=now)
|
|
if rotated is None:
|
|
return current.access_token, False
|
|
logger.info("Honcho OAuth token refreshed for host %s", host)
|
|
return rotated.access_token, True
|
|
|
|
|
|
def force_refresh_token(path: Path, host: str) -> str | None:
|
|
"""Rotate ``host``'s token now, ignoring local expiry (recovers a 401 on a
|
|
token the local clock still thinks is valid)."""
|
|
now = time.time()
|
|
key = (str(path), host)
|
|
with _refresh_lock, _config_refresh_lock(path):
|
|
cred = _load_cred(path, host)
|
|
if cred is None:
|
|
_expiry_cache.pop(key, None)
|
|
return None
|
|
# Dead grant, or an exchange just failed transiently: callers fail open.
|
|
if _grant_is_dead(key, cred) or _in_failure_cooldown(key):
|
|
return None
|
|
cached = _expiry_cache.get(key)
|
|
# Another thread or process already rotated: adopt the newer on-disk token.
|
|
if cached is not None and cred.access_token != cached[1] and not cred.is_expired(now=now):
|
|
_expiry_cache[key] = (cred.expires_at, cred.access_token)
|
|
return cred.access_token
|
|
rotated = _rotate_and_persist(path, host, key, cred, now=now, op_label="forced refresh")
|
|
if rotated is None:
|
|
return None
|
|
logger.info("Honcho OAuth token force-refreshed for host %s after an auth failure", host)
|
|
return rotated.access_token
|
|
|
|
|
|
def install_grant(
|
|
path: Path,
|
|
host: str,
|
|
grant: dict[str, Any],
|
|
*,
|
|
client_id: str,
|
|
token_endpoint: str,
|
|
apply_config: bool = True,
|
|
now: float | None = None,
|
|
) -> OAuthCredential:
|
|
"""Apply a fresh OAuth grant (an OAuthTokenResponse dict) to ``path`` for ``host``.
|
|
Deep-merges the grant's ``config`` into the file root (preserving other hosts
|
|
and root keys), then writes the host's ``apiKey`` and ``oauth`` block.
|
|
``apply_config=False`` stores tokens only."""
|
|
now = time.time() if now is None else now
|
|
cred = OAuthCredential.from_token_response(grant, now=now, client_id=client_id, token_endpoint=token_endpoint)
|
|
raw = _read_config(path)
|
|
granted_config = grant.get("config")
|
|
if isinstance(granted_config, dict):
|
|
cred.consent_peer_name = granted_config.get("peerName")
|
|
if apply_config:
|
|
_deep_merge(raw, granted_config)
|
|
_persist_credential(path, host, cred, raw)
|
|
return cred
|
|
|
|
|
|
def apply_token_to_client(client: Any, token: str) -> bool:
|
|
"""Rotate the live Honcho client's Bearer in place. The SDK builds its auth
|
|
header per request from ``_http.api_key``, so mutating it rotates every
|
|
holder of the singleton. Returns False on an SDK shape change (caller resets)."""
|
|
http = getattr(client, "_http", None)
|
|
if http is None or not hasattr(http, "api_key"):
|
|
return False
|
|
http.api_key = token
|
|
return True
|