mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 16:13:15 +00:00
2070a770c2
* fix: stop storing GitHub tokens in metadata Co-authored-by: open-swe[bot] <215916821+open-swe[bot]@users.noreply.github.com> * fix: bound in-process GitHub token cache with 24h TTL + sweep --------- Co-authored-by: open-swe[bot] <215916821+open-swe[bot]@users.noreply.github.com> Co-authored-by: open-swe[bot] <johannes@langchain.dev>
125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
"""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)
|
|
_GITHUB_TOKEN_CACHE: dict[str, tuple[str, str | None, datetime]] = {}
|
|
|
|
|
|
class GitHubAuthError(Exception):
|
|
"""Raised when a GitHub call returns 401, signalling a stale/revoked token."""
|
|
|
|
|
|
def cache_github_token_for_thread(
|
|
thread_id: str, token: str, expires_at: str | None = 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)
|
|
_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) 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) -> 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 = 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
|
|
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 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))
|
|
return token
|
|
|
|
|
|
async def get_github_token_from_thread(thread_id: str) -> tuple[str | None, str | None]:
|
|
"""Resolve the current process's cached GitHub token for a thread."""
|
|
return _cached_token_if_fresh(thread_id)
|
|
|
|
|
|
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)
|