from __future__ import annotations import time from typing import Any from yuxi.utils.logging_config import logger TOKEN_CACHE_TTL_S = 3600 TOKEN_REFRESH_MARGIN_S = 300 DEFAULT_FEISHU_DOMAIN = "https://open.feishu.cn" DEFAULT_LARK_DOMAIN = "https://open.larksuite.com" class TokenBackoffManager: """共享的 401/token 退避重试管理器""" def __init__(self, max_retries: int = 3, base_delay: float = 1.0, max_delay: float = 30.0): self._backoff_until: float = 0 self._retry_count: int = 0 self._max_retries = max_retries self._base_delay = base_delay self._max_delay = max_delay @property def is_backing_off(self) -> bool: return self._backoff_until > time.monotonic() async def wait_if_backing_off(self) -> None: import asyncio now = time.monotonic() if self._backoff_until > now: wait = self._backoff_until - now logger.warning(f"[TokenBackoff] Backoff active, waiting {wait:.1f}s") await asyncio.sleep(wait) def record_error(self) -> bool: self._retry_count += 1 if self._retry_count > self._max_retries: logger.error(f"[TokenBackoff] Retry count exceeded ({self._max_retries})") return False delay = min(self._base_delay * (2 ** (self._retry_count - 1)), self._max_delay) self._backoff_until = time.monotonic() + delay logger.warning(f"[TokenBackoff] Error recorded, retrying in {delay:.1f}s (attempt {self._retry_count})") return True def reset(self) -> None: if self._retry_count > 0: logger.info("[TokenBackoff] Backoff reset after success") self._retry_count = 0 self._backoff_until = 0 class FeishuOAuthClient: def __init__(self, app_id: str, app_secret: str, redirect_uri: str = "", domain: str = ""): self._app_id = app_id self._app_secret = app_secret self._redirect_uri = redirect_uri self._domain = domain or DEFAULT_FEISHU_DOMAIN self._token_cache: dict[str, tuple[dict[str, Any], float]] = {} self._access_token: str = "" self._access_token_expire_at: float = 0 self._refresh_token: str = "" def get_authorization_url(self, state: str = "", scope: str = "") -> str: base_url = f"{self._domain}/open-apis/authen/v1/index" params = { "app_id": self._app_id, "redirect_uri": self._redirect_uri or "http://localhost/callback", } if state: params["state"] = state if scope: params["scope"] = scope query = "&".join(f"{k}={v}" for k, v in params.items()) return f"{base_url}?{query}" async def exchange_code_for_token(self, code: str) -> dict[str, Any]: import httpx cache_key = f"code:{code}" cached = self._token_cache.get(cache_key) if cached is not None: token_data, cached_at = cached if time.monotonic() - cached_at < 30: return token_data del self._token_cache[cache_key] url = f"{self._domain}/open-apis/authen/v1/oidc/access_token" headers = {"Content-Type": "application/json"} body = { "app_id": self._app_id, "app_secret": self._app_secret, "grant_type": "authorization_code", "code": code, } async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client: resp = await client.post(url, headers=headers, json=body) if resp.status_code == 200: data = resp.json() self._token_cache[cache_key] = (data, time.monotonic()) if data.get("access_token"): self._access_token = data["access_token"] expires_in = data.get("expires_in", 7200) self._access_token_expire_at = time.monotonic() + expires_in - TOKEN_REFRESH_MARGIN_S if data.get("refresh_token"): self._refresh_token = data["refresh_token"] return data logger.error("[FeishuOAuth] Token exchange failed: HTTP %d", resp.status_code) return {} async def refresh_user_access_token(self, refresh_token: str = "") -> dict[str, Any]: import httpx rt = refresh_token or self._refresh_token if not rt: logger.error("[FeishuOAuth] No refresh token available") return {} cache_key = f"refresh:{rt}" cached = self._token_cache.get(cache_key) if cached is not None: token_data, cached_at = cached if time.monotonic() - cached_at < 30: return token_data del self._token_cache[cache_key] url = f"{self._domain}/open-apis/authen/v1/oidc/refresh_access_token" headers = {"Content-Type": "application/json"} body = { "app_id": self._app_id, "app_secret": self._app_secret, "grant_type": "refresh_token", "refresh_token": rt, } async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client: resp = await client.post(url, headers=headers, json=body) if resp.status_code == 200: data = resp.json() self._token_cache[cache_key] = (data, time.monotonic()) if data.get("access_token"): self._access_token = data["access_token"] expires_in = data.get("expires_in", 7200) self._access_token_expire_at = time.monotonic() + expires_in - TOKEN_REFRESH_MARGIN_S if data.get("refresh_token"): self._refresh_token = data["refresh_token"] return data logger.error("[FeishuOAuth] Refresh failed: HTTP %d", resp.status_code) return {} def is_access_token_expired(self) -> bool: if not self._access_token: return True return time.monotonic() >= self._access_token_expire_at def clear_cache(self) -> None: self._token_cache.clear() self._access_token = "" self._access_token_expire_at = 0 self._refresh_token = "" class FeishuDeviceCodeClient: def __init__(self, app_id: str, app_secret: str, domain: str = ""): self._app_id = app_id self._app_secret = app_secret self._domain = domain or DEFAULT_FEISHU_DOMAIN self._device_code_url = f"{self._domain}/open-apis/authen/v1/device/code" self._device_token_url = f"{self._domain}/open-apis/authen/v1/oidc/access_token" self._token_cache: dict[str, tuple[dict[str, Any], float]] = {} def init_device_flow(self) -> dict[str, Any] | None: import httpx body = { "app_id": self._app_id, "app_secret": self._app_secret, "scope": "user:read", } try: resp = httpx.post( self._device_code_url, json=body, headers={"Content-Type": "application/json"}, timeout=httpx.Timeout(30), ) if resp.status_code == 200: data = resp.json() if data.get("code") == 0: return data.get("data", {}) logger.error("[FeishuDeviceCode] init failed: code=%s, msg=%s", data.get("code"), data.get("msg")) else: logger.error("[FeishuDeviceCode] init HTTP %d: %s", resp.status_code, resp.text[:300]) except Exception as e: logger.error("[FeishuDeviceCode] init error: %s", e) return None def poll_device_token(self, device_code: str) -> dict[str, Any] | None: import httpx cache_key = f"device:{device_code}" cached = self._token_cache.get(cache_key) if cached is not None: token_data, cached_at = cached if time.monotonic() - cached_at < 5: return token_data del self._token_cache[cache_key] body = { "app_id": self._app_id, "app_secret": self._app_secret, "grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code, } try: resp = httpx.post( self._device_token_url, json=body, headers={"Content-Type": "application/json"}, timeout=httpx.Timeout(30), ) if resp.status_code == 200: data = resp.json() self._token_cache[cache_key] = (data, time.monotonic()) return data logger.error("[FeishuDeviceCode] poll HTTP %d: %s", resp.status_code, resp.text[:300]) except Exception as e: logger.error("[FeishuDeviceCode] poll error: %s", e) return None def clear_cache(self) -> None: self._token_cache.clear()