open-swe/agent/utils/github_token.py

"""GitHub token lookup utilities."""

from __future__ import annotations

import logging
from collections.abc import Mapping
from datetime import UTC, datetime, timedelta
from typing import Any

from langgraph.config import get_config

logger = logging.getLogger(__name__)

# Treat tokens with <= this many seconds remaining as expired so we re-auth
# before kicking off long agent runs.
_GITHUB_TOKEN_EXPIRY_SKEW_SECONDS = 60
# Hard cap on how long an entry stays cached regardless of the token's own
# expiry, so entries for threads that are never read again don't accumulate.
_GITHUB_TOKEN_MAX_TTL = timedelta(hours=24)
# thread_id -> (token, token_expires_at, cached_at, bound_repo). ``bound_repo``
# ("owner/name") binds the entry to the repo it was resolved for so a colliding
# thread_id originating from a different repo cannot reuse another repo's token.
_GITHUB_TOKEN_CACHE: dict[str, tuple[str, str | None, datetime, str | None]] = {}


class GitHubAuthError(Exception):
    """Raised when a GitHub call returns 401, signalling a stale/revoked token."""


def repo_cache_key(repo: Any) -> str | None:
    """Normalize a repo dict/string to a casefolded ``owner/name`` (None if unknown).

    GitHub owner/name are case-insensitive, so the key is casefolded on both the
    write (binding) and read (compare) sides to keep ``Org/Repo`` and ``org/repo``
    a single repo and avoid spurious cross-repo mismatches.
    """
    if isinstance(repo, str):
        cleaned = repo.strip()
        return cleaned.casefold() or None
    if isinstance(repo, Mapping):
        owner = repo.get("owner")
        name = repo.get("name")
        if isinstance(owner, str) and isinstance(name, str) and owner and name:
            return f"{owner}/{name}".casefold()
    return None


def cache_github_token_for_thread(
    thread_id: str, token: str, expires_at: str | None = None, *, repo: Any = None
) -> None:
    """Cache a GitHub token in process for the current thread."""
    if not thread_id or not token:
        return
    now = datetime.now(UTC)
    _GITHUB_TOKEN_CACHE[thread_id] = (token, expires_at, now, repo_cache_key(repo))
    _evict_expired(now=now)


def _is_expired(expires_at: Any, *, now: datetime | None = None) -> bool:
    """Return True when ``expires_at`` is past (or close to) ``now``."""
    if expires_at is None:
        return False

    parsed: datetime | None = None
    if isinstance(expires_at, int | float):
        try:
            parsed = datetime.fromtimestamp(float(expires_at), tz=UTC)
        except (OverflowError, OSError, ValueError):
            return False
    elif isinstance(expires_at, str):
        raw = expires_at.strip()
        if not raw:
            return False
        if raw.endswith("Z"):
            raw = raw[:-1] + "+00:00"
        try:
            parsed = datetime.fromisoformat(raw)
        except ValueError:
            return False

    if parsed is None:
        return False
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=UTC)

    current = (now or datetime.now(UTC)).astimezone(UTC)
    return (parsed - current).total_seconds() <= _GITHUB_TOKEN_EXPIRY_SKEW_SECONDS


def _entry_expired(expires_at: str | None, cached_at: datetime, *, now: datetime) -> bool:
    """Expired when past the token's own expiry or the 24h cache cap."""
    if now - cached_at >= _GITHUB_TOKEN_MAX_TTL:
        return True
    return _is_expired(expires_at, now=now)


def _evict_expired(*, now: datetime | None = None) -> None:
    current = now or datetime.now(UTC)
    stale = [
        tid
        for tid, (_token, expires_at, cached_at, _repo) in _GITHUB_TOKEN_CACHE.items()
        if _entry_expired(expires_at, cached_at, now=current)
    ]
    for tid in stale:
        _GITHUB_TOKEN_CACHE.pop(tid, None)


def _cached_token_if_fresh(
    thread_id: str | None, *, expected_repo: Any = None
) -> tuple[str | None, str | None]:
    if not thread_id:
        return None, None
    cached = _GITHUB_TOKEN_CACHE.get(thread_id)
    if not cached:
        return None, None
    token, expires_at, cached_at, bound_repo = cached
    if _entry_expired(expires_at, cached_at, now=datetime.now(UTC)):
        _GITHUB_TOKEN_CACHE.pop(thread_id, None)
        logger.info("Cached GitHub token for thread %s has expired; re-resolving", thread_id)
        return None, None
    expected = repo_cache_key(expected_repo)
    if expected and bound_repo and expected != bound_repo:
        _GITHUB_TOKEN_CACHE.pop(thread_id, None)
        logger.warning(
            "Cached GitHub token for thread %s is bound to repo %s, not %s; "
            "refusing cross-repo reuse",
            thread_id,
            bound_repo,
            expected,
        )
        return None, None
    return token, expires_at


def _thread_id_from_config(run_config: Mapping[str, Any]) -> str | None:
    configurable = run_config.get("configurable", {})
    if not isinstance(configurable, Mapping):
        return None
    thread_id = configurable.get("thread_id")
    return thread_id if isinstance(thread_id, str) and thread_id else None


def _repo_from_config(run_config: Mapping[str, Any]) -> Any:
    configurable = run_config.get("configurable", {})
    if not isinstance(configurable, Mapping):
        return None
    return configurable.get("repo")


def get_github_token(run_config: Mapping[str, Any] | None = None) -> str | None:
    """Resolve the current thread's GitHub token from process memory."""
    resolved = run_config if run_config is not None else get_config()
    token, _expires_at = _cached_token_if_fresh(
        _thread_id_from_config(resolved), expected_repo=_repo_from_config(resolved)
    )
    return token


async def get_github_token_from_thread(
    thread_id: str, *, expected_repo: Any = None
) -> tuple[str | None, str | None]:
    """Resolve the current process's cached GitHub token for a thread."""
    return _cached_token_if_fresh(thread_id, expected_repo=expected_repo)


async def invalidate_cached_github_token(thread_id: str) -> None:
    """Clear a cached GitHub token for a thread."""
    _GITHUB_TOKEN_CACHE.pop(thread_id, None)
    logger.info("Invalidated cached GitHub token for thread %s", thread_id)