"""Per-user third-party service credentials.""" from __future__ import annotations import asyncio import logging from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any from langgraph_sdk import get_client from pydantic import BaseModel, field_validator from ..encryption import decrypt_token, encrypt_token from .notion_oauth import is_reauth_required_error, refresh_notion_access_token logger = logging.getLogger(__name__) USER_CREDENTIALS_NAMESPACE: list[str] = ["user_credentials"] CURRENTS_KEY = "currents" NOTION_KEY = "notion" CURRENTS_API_BASE = "https://api.currents.dev/v1" _NOTION_TOKEN_EXPIRY_SKEW_SECONDS = 300 def _client(): return get_client() def _last4(value: str) -> str: return value[-4:] if len(value) >= 4 else value async def _get_provider(login: str, key: str) -> dict[str, Any] | None: try: item = await _client().store.get_item([*USER_CREDENTIALS_NAMESPACE, login], key) except Exception as e: # noqa: BLE001 logger.debug("user credentials lookup failed for %s/%s: %s", login, key, e) return None 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 _put_provider(login: str, key: str, value: dict[str, Any]) -> None: await _client().store.put_item([*USER_CREDENTIALS_NAMESPACE, login], key, value) async def _delete_provider(login: str, key: str) -> None: try: await _client().store.delete_item([*USER_CREDENTIALS_NAMESPACE, login], key) except Exception as e: # noqa: BLE001 logger.debug("user credentials delete failed for %s/%s: %s", login, key, e) class CurrentsCredentialsUpdate(BaseModel): """Connect Currents.dev with an organization API key.""" api_key: str @field_validator("api_key") @classmethod def _require_non_empty(cls, v: object) -> str: if not isinstance(v, str) or not v.strip(): raise ValueError("api_key must be a non-empty string") return v.strip() @dataclass(frozen=True) class NotionCredentials: access_token: str refresh_token: str | None token_endpoint: str client_id: str client_secret: str | None = None def _expires_at_from_response(data: dict[str, Any], *, field: str = "expires_in") -> str | None: raw = data.get(field) if not isinstance(raw, int | float) or raw <= 0: return None return (datetime.now(UTC) + timedelta(seconds=int(raw))).isoformat() def _token_expired( expires_at: str | None, *, skew_seconds: int = _NOTION_TOKEN_EXPIRY_SKEW_SECONDS ) -> 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 get_notion_status(login: str) -> dict[str, Any]: """Return a redacted view of the user's Notion MCP connection.""" notion = await _get_provider(login, NOTION_KEY) return { "notion": { "connected": True, "token_expires_at": notion.get("token_expires_at"), "updated_at": notion.get("updated_at"), } if notion else {"connected": False}, } def _notion_record_from_response( data: dict[str, Any], *, existing: dict[str, Any] | None = None, client_id: str | None = None, client_secret: str | None = None, token_endpoint: str | None = None, ) -> dict[str, Any]: existing = existing or {} access_token = data.get("access_token") if not isinstance(access_token, str) or not access_token: raise ValueError("Notion OAuth response missing access_token") refresh_token = data.get("refresh_token") record: dict[str, Any] = { "encrypted_access_token": encrypt_token(access_token), "client_id": client_id or existing.get("client_id", ""), "token_endpoint": token_endpoint or existing.get("token_endpoint", ""), "updated_at": datetime.now(UTC).isoformat(), } token_type = data.get("token_type") if isinstance(token_type, str): record["token_type"] = token_type scope = data.get("scope") if isinstance(scope, str): record["scope"] = scope token_expires_at = _expires_at_from_response(data) if token_expires_at: record["token_expires_at"] = token_expires_at elif existing.get("token_expires_at"): record["token_expires_at"] = existing["token_expires_at"] refresh_token_expires_at = _expires_at_from_response(data, field="refresh_token_expires_in") if refresh_token_expires_at: record["refresh_token_expires_at"] = refresh_token_expires_at elif existing.get("refresh_token_expires_at"): record["refresh_token_expires_at"] = existing["refresh_token_expires_at"] if isinstance(refresh_token, str) and refresh_token: record["encrypted_refresh_token"] = encrypt_token(refresh_token) elif existing.get("encrypted_refresh_token"): record["encrypted_refresh_token"] = existing["encrypted_refresh_token"] if client_secret: record["encrypted_client_secret"] = encrypt_token(client_secret) elif existing.get("encrypted_client_secret"): record["encrypted_client_secret"] = existing["encrypted_client_secret"] return record async def connect_notion(login: str, data: dict[str, Any], flow: dict[str, Any]) -> dict[str, Any]: client_id = flow.get("client_id") token_endpoint = flow.get("token_endpoint") if not isinstance(client_id, str) or not isinstance(token_endpoint, str): raise ValueError("stored Notion OAuth flow is incomplete") client_secret = ( flow.get("client_secret") if isinstance(flow.get("client_secret"), str) else None ) await _put_provider( login, NOTION_KEY, _notion_record_from_response( data, client_id=client_id, client_secret=client_secret, token_endpoint=token_endpoint, ), ) return await get_notion_status(login) async def disconnect_notion(login: str) -> dict[str, Any]: await _delete_provider(login, NOTION_KEY) return await get_notion_status(login) def _decrypt_notion_access_token(record: dict[str, Any]) -> str | None: token = decrypt_token(record.get("encrypted_access_token", "")) return token or None def _decrypt_notion_refresh_token(record: dict[str, Any]) -> str | None: token = decrypt_token(record.get("encrypted_refresh_token", "")) return token or None def _decrypt_notion_client_secret(record: dict[str, Any]) -> str | None: token = decrypt_token(record.get("encrypted_client_secret", "")) return token or None _notion_refresh_locks: dict[str, asyncio.Lock] = {} def _notion_refresh_lock(login: str) -> asyncio.Lock: lock = _notion_refresh_locks.get(login) if lock is None: lock = asyncio.Lock() _notion_refresh_locks[login] = lock return lock async def _refresh_stored_notion_token( login: str, record: dict[str, Any], ) -> tuple[str | None, bool]: refresh_token = _decrypt_notion_refresh_token(record) token_endpoint = record.get("token_endpoint") client_id = record.get("client_id") if not refresh_token or not isinstance(token_endpoint, str) or not isinstance(client_id, str): return None, False try: data = await refresh_notion_access_token( refresh_token=refresh_token, token_endpoint=token_endpoint, client_id=client_id, client_secret=_decrypt_notion_client_secret(record), ) except Exception as exc: # noqa: BLE001 logger.warning("Notion token refresh failed for %s", login, exc_info=True) return None, is_reauth_required_error(exc) await _put_provider(login, NOTION_KEY, _notion_record_from_response(data, existing=record)) access_token = data.get("access_token") return (access_token if isinstance(access_token, str) else None), False async def get_notion_credentials( login: str, *, force_refresh: bool = False ) -> NotionCredentials | None: """Return a valid Notion MCP credential set for a user.""" record = await _get_provider(login, NOTION_KEY) if not record: return None access_token = _decrypt_notion_access_token(record) if not access_token: return None if not force_refresh and not _token_expired(record.get("token_expires_at")): return NotionCredentials( access_token=access_token, refresh_token=_decrypt_notion_refresh_token(record), token_endpoint=record.get("token_endpoint", ""), client_id=record.get("client_id", ""), client_secret=_decrypt_notion_client_secret(record), ) if not _decrypt_notion_refresh_token(record): return None async with _notion_refresh_lock(login): record = await _get_provider(login, NOTION_KEY) if not record: return None access_token = _decrypt_notion_access_token(record) if not access_token: return None if not force_refresh and not _token_expired(record.get("token_expires_at")): return NotionCredentials( access_token=access_token, refresh_token=_decrypt_notion_refresh_token(record), token_endpoint=record.get("token_endpoint", ""), client_id=record.get("client_id", ""), client_secret=_decrypt_notion_client_secret(record), ) refreshed, refresh_token_dead = await _refresh_stored_notion_token(login, record) if refreshed: refreshed_record = await _get_provider(login, NOTION_KEY) or record return NotionCredentials( access_token=refreshed, refresh_token=_decrypt_notion_refresh_token(refreshed_record), token_endpoint=refreshed_record.get("token_endpoint", ""), client_id=refreshed_record.get("client_id", ""), client_secret=_decrypt_notion_client_secret(refreshed_record), ) if refresh_token_dead: latest = await _get_provider(login, NOTION_KEY) if latest and latest.get("encrypted_refresh_token") != record.get( "encrypted_refresh_token" ): latest_access_token = _decrypt_notion_access_token(latest) if latest_access_token: return NotionCredentials( access_token=latest_access_token, refresh_token=_decrypt_notion_refresh_token(latest), token_endpoint=latest.get("token_endpoint", ""), client_id=latest.get("client_id", ""), client_secret=_decrypt_notion_client_secret(latest), ) logger.info("Dropping dead Notion authorization for %s; reconnect required", login) await disconnect_notion(login) return None return NotionCredentials( access_token=access_token, refresh_token=_decrypt_notion_refresh_token(record), token_endpoint=record.get("token_endpoint", ""), client_id=record.get("client_id", ""), client_secret=_decrypt_notion_client_secret(record), ) async def get_notion_access_token(login: str) -> str | None: credentials = await get_notion_credentials(login) return credentials.access_token if credentials else None async def get_currents_status(login: str) -> dict[str, Any]: """Return a redacted, dashboard-safe view of the user's Currents key.""" currents = await _get_provider(login, CURRENTS_KEY) return { "currents": { "connected": True, "api_key_last4": currents.get("api_key_last4", ""), "updated_at": currents.get("updated_at"), } if currents else {"connected": False}, } async def connect_currents(login: str, update: CurrentsCredentialsUpdate) -> dict[str, Any]: await _put_provider( login, CURRENTS_KEY, { "encrypted_api_key": encrypt_token(update.api_key), "api_key_last4": _last4(update.api_key), "updated_at": datetime.now(UTC).isoformat(), }, ) return await get_currents_status(login) async def disconnect_currents(login: str) -> dict[str, Any]: await _delete_provider(login, CURRENTS_KEY) return await get_currents_status(login) async def get_currents_api_key(login: str) -> str | None: """Return the decrypted Currents API key, or ``None`` when not connected.""" currents = await _get_provider(login, CURRENTS_KEY) if not isinstance(currents, dict): return None api_key = decrypt_token(currents.get("encrypted_api_key", "")) return api_key or None