diff --git a/roboco/llm/providers/grok_auth.py b/roboco/llm/providers/grok_auth.py index 49d37bfc..201bf8e7 100644 --- a/roboco/llm/providers/grok_auth.py +++ b/roboco/llm/providers/grok_auth.py @@ -19,6 +19,7 @@ rather than hanging at the login prompt. from __future__ import annotations +import base64 import contextlib import json import os @@ -161,6 +162,42 @@ def _is_stale(creds: dict[str, Any], now: datetime, skew_seconds: int) -> bool: return expires_at is None or (expires_at - now).total_seconds() <= skew_seconds +# Documented grok access-token TTL (~6h, baked into the JWT ``exp`` by xAI's auth +# server). Used only as the last-resort ``expires_at`` when a refresh response +# omits ``expires_in`` AND the access token's JWT ``exp`` is unreadable. +_DEFAULT_ACCESS_TOKEN_TTL_HOURS = 6 +# A JWT is three ``.``-separated parts (header.payload.signature); the payload +# at index 1 carries the ``exp`` claim. Fewer parts means it isn't a JWT. +_MIN_JWT_PARTS = 2 + + +def _exp_from_access_token(access_token: str) -> datetime | None: + """Decode the JWT ``exp`` claim (unix seconds) from the access token. + + The grok access token is a JWT whose ``exp`` is the authoritative expiry. xAI + sometimes omits ``expires_in`` from the refresh response; without reading + ``exp`` the new token would keep the stale pre-refresh ``expires_at`` and + ``is_valid`` / ``--check`` would forever reject a fresh token. Returns + ``None`` when the token isn't a parseable JWT with a numeric ``exp``. + """ + # A JWT is ``header.payload.signature``; the payload (parts[1]) holds ``exp``. + parts = access_token.split(".") + if len(parts) < _MIN_JWT_PARTS: + return None + payload_b64 = parts[1] + padding = "=" * (-len(payload_b64) % 4) + try: + payload = json.loads(base64.urlsafe_b64decode(payload_b64 + padding)) + except (ValueError, json.JSONDecodeError): + return None + if not isinstance(payload, dict): + return None + exp = payload.get("exp") + if not isinstance(exp, (int, float)): + return None + return datetime.fromtimestamp(float(exp), tz=UTC) + + def _apply_refreshed_token( creds: dict[str, Any], token: dict[str, Any], now: datetime ) -> None: @@ -171,6 +208,26 @@ def _apply_refreshed_token( expires_in = token.get("expires_in") if isinstance(expires_in, (int, float)): creds["expires_at"] = _to_iso_z(now + timedelta(seconds=float(expires_in))) + else: + # xAI's refresh response sometimes omits expires_in. The access token + # is a JWT whose `exp` is the authoritative expiry — decode it so a fresh + # token isn't left with the stale pre-refresh expires_at (which would + # make is_valid / --check forever reject it AND make the refresh loop + # re-rotate the single-use refresh token every tick, killing the + # credential — see F006). + exp = _exp_from_access_token(token["access_token"]) + if exp is not None: + creds["expires_at"] = _to_iso_z(exp) + else: + creds["expires_at"] = _to_iso_z( + now + timedelta(hours=_DEFAULT_ACCESS_TOKEN_TTL_HOURS) + ) + logger.warning( + "grok auth refresh omitted expires_in and the access token's " + "JWT exp was unreadable; defaulting expires_at to %dh — " + "re-check on next tick", + _DEFAULT_ACCESS_TOKEN_TTL_HOURS, + ) creds["create_time"] = _to_iso_z(now) diff --git a/tests/unit/llm/providers/test_grok_auth.py b/tests/unit/llm/providers/test_grok_auth.py index f3c9d92f..c2e9da53 100644 --- a/tests/unit/llm/providers/test_grok_auth.py +++ b/tests/unit/llm/providers/test_grok_auth.py @@ -2,9 +2,10 @@ from __future__ import annotations +import base64 import json import pathlib -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any from roboco.llm.providers import grok_auth as ga @@ -19,6 +20,17 @@ _FUTURE = "2099-01-01T00:00:00.000000000Z" _CLIENT = "b1a00492-client" +def _jwt(exp_unix: int) -> str: + """Build a minimal JWT (header.payload.signature) carrying an ``exp`` claim.""" + payload = ( + base64.urlsafe_b64encode(json.dumps({"exp": exp_unix}).encode()) + .rstrip(b"=") + .decode() + ) + header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode() + return f"{header}.{payload}.sig" + + def _bundle(expires_at: str, *, refresh_token: str = "rt") -> dict[str, Any]: return { f"https://auth.x.ai::{_CLIENT}": { @@ -99,6 +111,49 @@ def test_refresh_mints_new_token_when_stale(tmp_path: Path) -> None: assert ga.is_valid(path) # fresh expires_at ~6h out, valid against real now +def test_refresh_omitting_expires_in_still_marks_token_valid(tmp_path: Path) -> None: + """F092: if xAI's refresh response omits ``expires_in``, the access token's + JWT ``exp`` claim is the authoritative expiry — decode it so a fresh token + isn't left with the stale pre-refresh ``expires_at`` (which would make + ``is_valid`` / ``--check`` forever reject it and the refresh loop re-rotate + the single-use refresh token every tick).""" + path = tmp_path / "auth.json" + _write(path, _bundle(_PAST)) + exp_unix = int((datetime.now(UTC) + timedelta(hours=6)).timestamp()) + jwt_token = _jwt(exp_unix) + + def _post(_url: str, _form: dict[str, str]) -> dict[str, Any]: + # No expires_in — only the JWT access_token carries the expiry. + return {"access_token": jwt_token, "refresh_token": "new-rt"} + + assert ga.refresh_if_stale(path, post=_post) == "refreshed" + creds = next(iter(json.loads(path.read_text()).values())) + assert creds["key"] == jwt_token + assert creds["expires_at"] != _PAST # updated, not the stale pre-refresh value + assert ga.is_valid(path) # JWT exp ~6h out -> valid against real now + + +def test_refresh_omitting_expires_in_with_unreadable_jwt_defaults_ttl( + tmp_path: Path, +) -> None: + """F092 fallback: expires_in missing AND the access token isn't a JWT with a + readable ``exp`` — default to the documented ~6h TTL so a fresh token is + treated as live instead of stale, rather than forever rejected (the warning + is emitted via structlog, visible in the captured stdout).""" + path = tmp_path / "auth.json" + _write(path, _bundle(_PAST)) + + def _post(_url: str, _form: dict[str, str]) -> dict[str, Any]: + # No expires_in and a non-JWT access token (no `.`-separated payload). + return {"access_token": "opaque-not-a-jwt", "refresh_token": "new-rt"} + + assert ga.refresh_if_stale(path, post=_post) == "refreshed" + creds = next(iter(json.loads(path.read_text()).values())) + assert creds["key"] == "opaque-not-a-jwt" + assert creds["expires_at"] != _PAST # defaulted forward, not left stale + assert ga.is_valid(path) # ~6h default -> valid against real now + + def test_refresh_missing_file(tmp_path: Path) -> None: assert ga.refresh_if_stale(tmp_path / "nope.json") == "missing"