"""User profile schema and LangGraph Store CRUD. Storage is split into two namespaces to avoid the read-modify-write race between profile-edit writes and OAuth-callback token refreshes: * ``["profiles"]`` — user-editable settings (model, effort, default_repo). * ``["oauth_tokens"]`` — encrypted GitHub OAuth access token + email. Each upsert only touches its own namespace, so the two flows can't clobber each other's fields even when they interleave. """ from __future__ import annotations import asyncio import logging from datetime import UTC, datetime, timedelta from typing import Any import httpx from langgraph_sdk import get_client from pydantic import BaseModel, field_validator from ..encryption import decrypt_token, encrypt_token from .oauth import ( expires_at_from_github_response, is_unrecoverable_refresh_error, refresh_user_access_token, ) from .options import SUPPORTED_MODEL_IDS, model_supports_effort logger = logging.getLogger(__name__) PROFILES_NAMESPACE: list[str] = ["profiles"] OAUTH_TOKENS_NAMESPACE: list[str] = ["oauth_tokens"] class ProfileUpdate(BaseModel): default_model: str reasoning_effort: str default_subagent_model: str | None = None subagent_reasoning_effort: str | None = None default_repo: str | None = None base_branch: str | None = None branch_prefix: str | None = None auto_fix_ci: bool = True create_prs: bool = False review_draft_prs: bool | None = None @field_validator("default_model") @classmethod def _model_supported(cls, v: str) -> str: if v not in SUPPORTED_MODEL_IDS: raise ValueError(f"unsupported model: {v}") return v def validate_pairing(self) -> None: if not model_supports_effort(self.default_model, self.reasoning_effort): raise ValueError( f"effort {self.reasoning_effort!r} not supported by {self.default_model!r}" ) if self.default_subagent_model is None and self.subagent_reasoning_effort is None: return if self.default_subagent_model is None: raise ValueError("subagent reasoning effort set without a model") if self.default_subagent_model not in SUPPORTED_MODEL_IDS: raise ValueError(f"unsupported subagent model: {self.default_subagent_model}") if self.subagent_reasoning_effort is None or not model_supports_effort( self.default_subagent_model, self.subagent_reasoning_effort, ): raise ValueError( f"effort {self.subagent_reasoning_effort!r} not supported by " f"{self.default_subagent_model!r}" ) def _client(): return get_client() async def _get_value(namespace: list[str], key: str) -> dict[str, Any] | None: try: item = await _client().store.get_item(namespace, key) except httpx.HTTPStatusError as e: if e.response.status_code == 404: return None raise if item is None: return None value = item.get("value") if isinstance(item, dict) else getattr(item, "value", None) return value if isinstance(value, dict) else None async def get_profile(login: str) -> dict[str, Any] | None: return await _get_value(PROFILES_NAMESPACE, login) async def upsert_profile(login: str, email: str, update: ProfileUpdate) -> dict[str, Any]: """Write the user's editable settings. Only touches ``["profiles"]`` — the OAuth token in ``["oauth_tokens"]`` is untouched, so a concurrent re-login can't be clobbered by this write and vice versa. """ existing = await get_profile(login) or {} value: dict[str, Any] = { **existing, "login": login, "email": email or existing.get("email", ""), "default_model": update.default_model, "reasoning_effort": update.reasoning_effort, "default_subagent_model": update.default_subagent_model, "subagent_reasoning_effort": update.subagent_reasoning_effort, "default_repo": update.default_repo, "base_branch": update.base_branch, "branch_prefix": update.branch_prefix, "auto_fix_ci": update.auto_fix_ci, "create_prs": update.create_prs, "review_draft_prs": update.review_draft_prs, "updated_at": datetime.now(UTC).isoformat(), } for stale_field in ( "first_name", "last_name", "allow_artifacts", "slack_notifications", "preferred_pr_destination", ): value.pop(stale_field, None) await _client().store.put_item(PROFILES_NAMESPACE, login, value) return value _refresh_locks: dict[str, asyncio.Lock] = {} def _refresh_lock(login: str) -> asyncio.Lock: lock = _refresh_locks.get(login) if lock is None: lock = asyncio.Lock() _refresh_locks[login] = lock return lock def _token_expired(expires_at: str | None, *, skew_seconds: int = 300) -> bool: if not isinstance(expires_at, str) or not expires_at: return False try: exp = datetime.fromisoformat(expires_at.replace("Z", "+00:00")) if exp.tzinfo is None: exp = exp.replace(tzinfo=UTC) except ValueError: return False return datetime.now(UTC) + timedelta(seconds=skew_seconds) >= exp async def upsert_access_token( login: str, email: str, access_token: str, *, refresh_token: str | None = None, token_expires_at: str | None = None, refresh_token_expires_at: str | None = None, ) -> None: """Persist (or refresh) the user's encrypted GitHub OAuth tokens. Only touches ``["oauth_tokens"]`` — the user-editable profile is left intact even if a save is in flight in another request. """ if not access_token: return existing = await _get_value(OAUTH_TOKENS_NAMESPACE, login) or {} value: dict[str, Any] = { "login": login, "email": email or existing.get("email", ""), "encrypted_gh_token": encrypt_token(access_token), "updated_at": datetime.now(UTC).isoformat(), } if refresh_token: value["encrypted_gh_refresh_token"] = encrypt_token(refresh_token) elif existing.get("encrypted_gh_refresh_token"): value["encrypted_gh_refresh_token"] = existing["encrypted_gh_refresh_token"] if token_expires_at: value["token_expires_at"] = token_expires_at if refresh_token_expires_at: value["refresh_token_expires_at"] = refresh_token_expires_at await _client().store.put_item(OAUTH_TOKENS_NAMESPACE, login, value) async def upsert_access_token_from_github_response( login: str, email: str, data: dict[str, Any] ) -> None: """Store tokens from a GitHub OAuth code exchange or refresh response.""" access_token = data.get("access_token") if not isinstance(access_token, str) or not access_token: return refresh_token = data.get("refresh_token") await upsert_access_token( login, email, access_token, refresh_token=refresh_token if isinstance(refresh_token, str) else None, token_expires_at=expires_at_from_github_response(data, field="expires_in"), refresh_token_expires_at=expires_at_from_github_response( data, field="refresh_token_expires_in" ), ) async def delete_access_token(login: str) -> None: """Drop the user's stored OAuth tokens. Used when a refresh token is permanently dead so we stop handing out a known-stale access token and callers prompt a clean re-login instead. """ try: await _client().store.delete_item(OAUTH_TOKENS_NAMESPACE, login) except httpx.HTTPStatusError as exc: if exc.response.status_code != 404: raise def _decrypt_access_token(record: dict[str, Any]) -> str | None: encrypted = record.get("encrypted_gh_token") if not encrypted: return None return decrypt_token(encrypted) or None def _decrypt_refresh_token(record: dict[str, Any]) -> str | None: encrypted = record.get("encrypted_gh_refresh_token") if not encrypted: return None return decrypt_token(encrypted) or None async def _refresh_stored_token(login: str, record: dict[str, Any]) -> tuple[str | None, bool]: """Refresh the stored token, returning ``(access_token, refresh_token_dead)``. ``refresh_token_dead`` is True when GitHub says the refresh token can never mint a new token again, so the caller should drop the stored authorization rather than keep serving a stale access token. """ refresh_token = _decrypt_refresh_token(record) if not refresh_token: return None, False try: data = await refresh_user_access_token(refresh_token) except Exception as exc: # noqa: BLE001 logger.warning("GitHub token refresh failed for %s", login, exc_info=True) return None, is_unrecoverable_refresh_error(exc) email = record.get("email") if isinstance(record.get("email"), str) else "" await upsert_access_token_from_github_response(login, email, data) access_token = data.get("access_token") return (access_token if isinstance(access_token, str) else None), False async def get_valid_access_token(login: str, *, force_refresh: bool = False) -> str | None: """Return a GitHub access token, refreshing proactively when near expiry.""" record = await _get_value(OAUTH_TOKENS_NAMESPACE, login) if not record: return None access_token = _decrypt_access_token(record) if not access_token: return None if not force_refresh and not _token_expired(record.get("token_expires_at")): return access_token if not _decrypt_refresh_token(record): return access_token async with _refresh_lock(login): record = await _get_value(OAUTH_TOKENS_NAMESPACE, login) if not record: return None access_token = _decrypt_access_token(record) if not access_token: return None if not force_refresh and not _token_expired(record.get("token_expires_at")): return access_token refreshed, refresh_token_dead = await _refresh_stored_token(login, record) if refreshed: return refreshed if refresh_token_dead: # The refresh token is permanently invalid (revoked / expired), so # the cached access token is dead too. Drop it so callers prompt a # clean re-login instead of repeatedly handing out a stale token. # The OAuth callback can write a fresh authorization while the # refresh request is in flight (it doesn't take this lock), so only # delete if the stored record is still the one that failed. latest = await _get_value(OAUTH_TOKENS_NAMESPACE, login) if latest and latest.get("encrypted_gh_refresh_token") != record.get( "encrypted_gh_refresh_token" ): return _decrypt_access_token(latest) logger.info("Dropping dead GitHub authorization for %s; re-login required", login) await delete_access_token(login) return None return access_token async def get_access_token(login: str) -> str | None: return await get_valid_access_token(login) async def has_access_token_record(login: str) -> bool: """Whether an OAuth token record exists for ``login``. Distinguishes "user has never completed a GitHub login" (no record) from "the stored authorization is present but no longer usable" (record exists but won't decrypt / was revoked), so callers can prompt accurately. """ return bool(await _get_value(OAUTH_TOKENS_NAMESPACE, login)) async def list_profiles() -> list[dict[str, Any]]: result = await _client().store.search_items(PROFILES_NAMESPACE, limit=1000) items = result.get("items") if isinstance(result, dict) else getattr(result, "items", []) out: list[dict[str, Any]] = [] for item in items or []: value = item.get("value") if isinstance(item, dict) else getattr(item, "value", None) if isinstance(value, dict): out.append(value) return out