"""Reusable Microsoft Graph REST client helpers.""" from __future__ import annotations import asyncio import os from pathlib import Path from typing import Any, Awaitable, Callable import httpx from agent.retry_utils import parse_retry_after_seconds from tools.microsoft_graph_auth import MicrosoftGraphTokenProvider, format_graph_error DEFAULT_GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0" Headers = dict[str, str] | None Params = dict[str, Any] | None class MicrosoftGraphClientError(RuntimeError): """Base class for Graph client failures.""" class MicrosoftGraphAPIError(MicrosoftGraphClientError): """Raised when a Graph API request fails.""" def __init__(self, status_code: int, method: str, url: str, message: str, *, retry_after_seconds: float | None = None, payload: Any = None) -> None: self.status_code, self.method, self.url = status_code, method, url self.retry_after_seconds, self.payload = retry_after_seconds, payload super().__init__(f"Microsoft Graph API error {status_code} for {method} {url}: {message}") class MicrosoftGraphClient: """Minimal async Graph client. Retry policy (JSON requests and streaming downloads alike): transport errors back off exponentially; 401 clears the token cache and refetches; 429/5xx honor ``Retry-After``. Each attempt uses a fresh ``AsyncClient``.""" def __init__(self, token_provider: MicrosoftGraphTokenProvider, *, base_url: str = DEFAULT_GRAPH_BASE_URL, timeout: float = 60.0, max_retries: int = 3, transport: httpx.AsyncBaseTransport | None = None, sleep: Callable[[float], Awaitable[None]] | None = None, user_agent: str = "Hermes-Agent/graph-client") -> None: self.token_provider, self.base_url, self.timeout = token_provider, base_url.rstrip("/"), timeout self.max_retries, self.user_agent = max(0, int(max_retries)), user_agent self._transport, self._sleep = transport, sleep or asyncio.sleep async def get_json(self, path: str, *, params: Params = None, headers: Headers = None) -> Any: return self._decode_json(await self._request("GET", path, params=params, headers=headers)) async def post_json(self, path: str, *, json_body: Any | None = None, headers: Headers = None) -> Any: return self._decode_json(await self._request("POST", path, json_body=json_body, headers=headers)) async def patch_json(self, path: str, *, json_body: Any | None = None, headers: Headers = None) -> Any: """Decoded body, or ``{}`` for a 204 / bodiless response.""" response = await self._request("PATCH", path, json_body=json_body, headers=headers) return self._decode_json(response) if response.status_code != 204 and response.content else {} async def delete(self, path: str, *, headers: Headers = None) -> dict[str, Any]: """Decoded body, or ``{"deleted": True, "status_code"}`` for a 204 / bodiless response.""" response = await self._request("DELETE", path, headers=headers) if response.status_code != 204 and response.content: return self._decode_json(response) return {"deleted": True, "status_code": response.status_code} async def collect_paginated(self, path: str, *, params: Params = None, headers: Headers = None) -> list[Any]: """Follow ``@odata.nextLink`` and concatenate every page's ``value`` list.""" items: list[Any] = [] # Query params go on the first request only; @odata.nextLink already embeds them. next_url, next_params = self._resolve_url(path), dict(params or {}) while next_url: payload = self._decode_json(await self._request("GET", next_url, params=next_params or None, headers=headers)) if not isinstance(payload, dict): raise MicrosoftGraphClientError( f"Expected paginated Graph response dict, got {type(payload).__name__}.") if isinstance(payload.get("value"), list): items.extend(payload["value"]) next_url, next_params = payload.get("@odata.nextLink"), {} return items async def download_to_file(self, path: str, destination: str | Path, *, headers: Headers = None, chunk_size: int = 65536) -> dict[str, Any]: """Stream a Graph resource to disk chunk-by-chunk (large recordings never fit in memory); written to ``.part`` and renamed into place only on success.""" url, target = self._resolve_url(path), Path(destination) target.parent.mkdir(parents=True, exist_ok=True) tmp_target = target.with_suffix(target.suffix + ".part") async def perform(client: httpx.AsyncClient, request_headers: dict[str, str]): try: async with client.stream("GET", url, headers=request_headers) as response: if response.status_code >= 400: await response.aread() # small error body -> meaningful message return response, None with tmp_target.open("wb") as handle: async for chunk in response.aiter_bytes(chunk_size=chunk_size): if chunk: handle.write(chunk) return response, response.headers.get("content-type") except httpx.HTTPError: tmp_target.unlink(missing_ok=True) raise content_type = await self._with_retries("GET", url, "*/*", None, headers, perform, "download") os.replace(tmp_target, target) return {"path": str(target), "size_bytes": target.stat().st_size, "content_type": content_type} async def _request(self, method: str, path_or_url: str, *, params: Params = None, json_body: Any | None = None, headers: Headers = None) -> httpx.Response: url = self._resolve_url(path_or_url) async def perform(client: httpx.AsyncClient, request_headers: dict[str, str]): response = await client.request(method, url, params=params, json=json_body, headers=request_headers) return response, response return await self._with_retries(method, url, "application/json", json_body, headers, perform, "request") async def _with_retries( self, method: str, url: str, accept: str, json_body: Any | None, headers: Headers, perform: Callable[[httpx.AsyncClient, dict[str, str]], Awaitable[tuple[httpx.Response, Any]]], kind: str) -> Any: """Run ``perform`` (-> ``(response, result)``) under the retry policy. ``kind`` only labels transport-failure messages. Raises ``MicrosoftGraphAPIError`` once retries are exhausted or the status is not retryable; only 401 forces a token refresh.""" last_error: Exception | None = None for attempt in range(self.max_retries + 1): token = await self.token_provider.get_access_token( force_refresh=isinstance(last_error, MicrosoftGraphAPIError) and last_error.status_code == 401) request_headers = {"Authorization": f"Bearer {token}", "Accept": accept, "User-Agent": self.user_agent, **({"Content-Type": "application/json"} if json_body is not None else {}), **(headers or {})} exhausted = attempt >= self.max_retries try: async with httpx.AsyncClient(timeout=httpx.Timeout(self.timeout), transport=self._transport) as client: response, result = await perform(client, request_headers) except httpx.HTTPError as exc: last_error, response = exc, None if exhausted: raise MicrosoftGraphClientError( f"Microsoft Graph {kind} failed for {method} {url}: {exc}") from exc else: if response.status_code < 400: return result last_error, status = self._build_api_error(method, url, response), response.status_code if exhausted or not (status in (401, 429) or 500 <= status < 600): raise last_error if status == 401: self.token_provider.clear_cache() await self._sleep(self._retry_delay(response, attempt)) raise MicrosoftGraphClientError(f"Microsoft Graph {kind} exhausted retries for {method} {url}.") def _resolve_url(self, path_or_url: str) -> str: if path_or_url.startswith(("http://", "https://")): return path_or_url return f"{self.base_url}{path_or_url if path_or_url.startswith('/') else '/' + path_or_url}" @staticmethod def _decode_json(response: httpx.Response) -> Any: try: return response.json() except ValueError as exc: raise MicrosoftGraphClientError( "Microsoft Graph response was not valid JSON for " f"{response.request.method} {response.request.url}") from exc @staticmethod def _retry_delay(response: httpx.Response | None, attempt: int) -> float: retry_after = parse_retry_after_seconds(response.headers) if response is not None else None return min(8.0, 0.5 * (2 ** attempt)) if retry_after is None else retry_after @staticmethod def _build_api_error(method: str, url: str, response: httpx.Response) -> MicrosoftGraphAPIError: try: payload: Any = response.json() except ValueError: payload = None detail = format_graph_error(payload.get("error")) if isinstance(payload, dict) else None return MicrosoftGraphAPIError( response.status_code, method, url, response.text.strip() or "unknown error" if detail is None else detail, retry_after_seconds=parse_retry_after_seconds(response.headers), payload=payload)