archastro.platform.runtime.http_client

  1# Runtime: async HTTP client for the generated Platform SDK.
  2# This file is hand-maintained, not generated.
  3
  4from __future__ import annotations
  5
  6import asyncio
  7import json
  8import threading
  9from collections.abc import AsyncIterator, Callable, Coroutine, Iterator
 10from functools import cache
 11from typing import Any, TypeVar, overload
 12
 13import httpx
 14from pydantic import TypeAdapter
 15
 16DEFAULT_API_PREFIX = "/api/v1"
 17
 18T = TypeVar("T")
 19
 20
 21@cache
 22def _type_adapter(tp: Any) -> TypeAdapter[Any]:
 23    """Cache adapters per response type; TypeAdapter construction is costly."""
 24    return TypeAdapter(tp)
 25
 26
 27class ApiError(Exception):
 28    """Structured API error with status code, error code, and message."""
 29
 30    def __init__(
 31        self,
 32        status: int,
 33        error_code: str,
 34        message: str,
 35        body: dict[str, Any] | None = None,
 36    ):
 37        super().__init__(message)
 38        self.status = status
 39        self.error_code = error_code
 40        self.body = body
 41
 42
 43class HttpClient:
 44    def __init__(
 45        self,
 46        *,
 47        base_url: str,
 48        access_token: str | None = None,
 49        get_access_token: Callable[[], str | None] | None = None,
 50        on_refresh_token: Callable[[], Coroutine[Any, Any, str]] | None = None,
 51        path_prefix: str | None = None,
 52        default_headers: dict[str, str] | None = None,
 53        refresh_only: bool = False,
 54    ):
 55        self._base_url = base_url.rstrip("/")
 56        self._access_token = access_token
 57        self._get_access_token = get_access_token
 58        self._on_refresh_token = on_refresh_token
 59        self._path_prefix = path_prefix
 60        self._default_headers = default_headers or {}
 61        self._client = httpx.AsyncClient(timeout=30.0)
 62        self._refresh_task: asyncio.Task[str] | None = None
 63        self._refresh_only = refresh_only
 64
 65    def _get_token(self) -> str | None:
 66        if self._get_access_token:
 67            return self._get_access_token()
 68        return self._access_token
 69
 70    def _transform_path(self, path: str) -> str:
 71        if self._path_prefix is None:
 72            return path
 73        if path.startswith(DEFAULT_API_PREFIX):
 74            return self._path_prefix + path[len(DEFAULT_API_PREFIX) :]
 75        return path
 76
 77    def set_access_token(self, token: str) -> None:
 78        self._access_token = token
 79
 80    def set_refresh_handler(self, handler: Callable[[], Coroutine[Any, Any, str]]) -> None:
 81        self._on_refresh_token = handler
 82
 83    async def _do_fetch(
 84        self,
 85        path: str,
 86        *,
 87        method: str = "GET",
 88        body: Any = None,
 89        headers: dict[str, str] | None = None,
 90        query: dict[str, Any] | None = None,
 91    ) -> httpx.Response:
 92        token = self._get_token()
 93        url = f"{self._base_url}{self._transform_path(path)}"
 94
 95        req_headers = {
 96            **self._default_headers,
 97            "Content-Type": "application/json",
 98        }
 99        if token:
100            req_headers["Authorization"] = f"Bearer {token}"
101        if headers:
102            req_headers.update(headers)
103
104        params = None
105        if query:
106            params = {k: v for k, v in query.items() if v is not None}
107
108        return await self._client.request(
109            method,
110            url,
111            json=body if body is not None and method not in ("GET", "HEAD") else None,
112            headers=req_headers,
113            params=params,
114        )
115
116    async def _execute(
117        self,
118        path: str,
119        *,
120        method: str = "GET",
121        body: Any = None,
122        headers: dict[str, str] | None = None,
123        query: dict[str, Any] | None = None,
124    ) -> httpx.Response:
125        """Fetch with auth gate, 401 auto-refresh, and error handling.
126
127        Returns the successful response for callers to interpret (JSON, raw bytes, etc.).
128        """
129        auth_prefix = f"{DEFAULT_API_PREFIX}/auth/"
130        if self._refresh_only and not path.startswith(auth_prefix):
131            raise RuntimeError(
132                f"Refresh-only HTTP client cannot make requests outside {auth_prefix}"
133            )
134
135        response = await self._do_fetch(
136            path, method=method, body=body, headers=headers, query=query
137        )
138
139        # Auto-refresh: on 401, attempt one token refresh and retry.
140        # The refresh handler runs on a separate HttpClient (refresh_only),
141        # so it cannot re-enter this block. Concurrent 401s piggyback on
142        # the same _refresh_task.
143        if (
144            response.status_code == 401
145            and self._on_refresh_token
146            and not path.startswith(auth_prefix)
147        ):
148            if self._refresh_task is None:
149
150                async def _do_refresh() -> str:
151                    try:
152                        return await self._on_refresh_token()  # type: ignore[misc]
153                    finally:
154                        self._refresh_task = None
155
156                self._refresh_task = asyncio.create_task(_do_refresh())
157            try:
158                new_token = await self._refresh_task
159            except Exception:
160                pass  # refresh failed — fall through to throw original 401
161            else:
162                self._access_token = new_token
163                response = await self._do_fetch(
164                    path, method=method, body=body, headers=headers, query=query
165                )
166
167        if response.status_code >= 400:
168            raw_data: dict[str, Any] = {}
169            try:
170                raw_data = response.json()
171            except Exception:
172                pass
173            error_code, message = _parse_error(raw_data, response.status_code)
174            raise ApiError(response.status_code, error_code, message, raw_data)
175
176        return response
177
178    @overload
179    async def request(
180        self,
181        path: str,
182        *,
183        method: str = "GET",
184        body: Any = None,
185        headers: dict[str, str] | None = None,
186        query: dict[str, Any] | None = None,
187        response_type: type[T],
188    ) -> T: ...
189
190    @overload
191    async def request(
192        self,
193        path: str,
194        *,
195        method: str = "GET",
196        body: Any = None,
197        headers: dict[str, str] | None = None,
198        query: dict[str, Any] | None = None,
199        response_type: None = None,
200    ) -> Any: ...
201
202    async def request(
203        self,
204        path: str,
205        *,
206        method: str = "GET",
207        body: Any = None,
208        headers: dict[str, str] | None = None,
209        query: dict[str, Any] | None = None,
210        response_type: type[T] | None = None,
211    ) -> T | Any:
212        """Issue a request and return the JSON body.
213
214        With `response_type`, the body is validated into that type (a Pydantic
215        model or a generic alias like list[Model]); a bodyless 204 then raises
216        ValidationError, since the operation promised a typed body. Without
217        `response_type`, the raw parsed JSON is returned, or None on 204.
218        """
219        response = await self._execute(path, method=method, body=body, headers=headers, query=query)
220
221        if response.status_code == 204:
222            if response_type is None:
223                return None
224            # A 204 on an operation that promises a typed body is a server
225            # contract violation; validating None fails loudly here instead
226            # of surfacing later as an AttributeError far from the call.
227            return _type_adapter(response_type).validate_python(None)
228
229        raw = response.json()
230        if response_type is None:
231            return raw
232        return _type_adapter(response_type).validate_python(raw)
233
234    async def request_raw(
235        self,
236        path: str,
237        *,
238        method: str = "GET",
239        body: Any = None,
240        headers: dict[str, str] | None = None,
241        query: dict[str, Any] | None = None,
242    ) -> dict[str, Any]:
243        response = await self._execute(path, method=method, body=body, headers=headers, query=query)
244
245        return {
246            "content": response.content,
247            "mime_type": response.headers.get("content-type", "text/plain"),
248        }
249
250    async def stream_sse(
251        self,
252        path: str,
253        *,
254        method: str = "GET",
255        body: Any = None,
256        headers: dict[str, str] | None = None,
257        query: dict[str, Any] | None = None,
258    ) -> AsyncIterator[dict[str, Any]]:
259        """Open a Server-Sent Events stream, yielding parsed ``{"event", "data"}``.
260
261        Backs the generated async ``stream()`` methods for ``x-sdk-streaming``
262        endpoints. Raises :class:`ApiError` on a non-2xx response before the
263        stream opens. (Token refresh is not retried mid-stream.)
264        """
265        auth_prefix = f"{DEFAULT_API_PREFIX}/auth/"
266        if self._refresh_only and not path.startswith(auth_prefix):
267            raise RuntimeError(
268                f"Refresh-only HTTP client cannot make requests outside {auth_prefix}"
269            )
270
271        url = f"{self._base_url}{self._transform_path(path)}"
272        sends_body = body is not None and method not in ("GET", "HEAD")
273        req_headers = {**self._default_headers, "Accept": "text/event-stream"}
274        if sends_body:
275            req_headers["Content-Type"] = "application/json"
276        token = self._get_token()
277        if token:
278            req_headers["Authorization"] = f"Bearer {token}"
279        if headers:
280            req_headers.update(headers)
281        params = {k: v for k, v in (query or {}).items() if v is not None} or None
282
283        async with self._client.stream(
284            method,
285            url,
286            json=body if sends_body else None,
287            headers=req_headers,
288            params=params,
289        ) as response:
290            if response.status_code >= 400:
291                await response.aread()
292                raw: dict[str, Any] = {}
293                try:
294                    raw = response.json()
295                except Exception:
296                    pass
297                code, message = _parse_error(raw, response.status_code)
298                raise ApiError(response.status_code, code, message, raw)
299
300            event: str | None = None
301            data_lines: list[str] = []
302            async for line in response.aiter_lines():
303                if line == "":
304                    parsed = _build_sse_event(event, data_lines)
305                    if parsed is not None:
306                        yield parsed
307                    event, data_lines = None, []
308                elif line.startswith("event:"):
309                    event = line[6:].strip()
310                elif line.startswith("data:"):
311                    data_lines.append(line[5:].strip())
312            parsed = _build_sse_event(event, data_lines)
313            if parsed is not None:
314                yield parsed
315
316    async def close(self) -> None:
317        await self._client.aclose()
318
319
320class SyncHttpClient:
321    def __init__(
322        self,
323        *,
324        base_url: str,
325        access_token: str | None = None,
326        get_access_token: Callable[[], str | None] | None = None,
327        on_refresh_token: Callable[[], str] | None = None,
328        path_prefix: str | None = None,
329        default_headers: dict[str, str] | None = None,
330        refresh_only: bool = False,
331    ):
332        self._base_url = base_url.rstrip("/")
333        self._access_token = access_token
334        self._get_access_token = get_access_token
335        self._on_refresh_token = on_refresh_token
336        self._path_prefix = path_prefix
337        self._default_headers = default_headers or {}
338        self._client = httpx.Client(timeout=30.0)
339        self._refresh_only = refresh_only
340        self._refresh_lock = threading.Lock()
341
342    def _get_token(self) -> str | None:
343        if self._get_access_token:
344            return self._get_access_token()
345        return self._access_token
346
347    def _transform_path(self, path: str) -> str:
348        if self._path_prefix is None:
349            return path
350        if path.startswith(DEFAULT_API_PREFIX):
351            return self._path_prefix + path[len(DEFAULT_API_PREFIX) :]
352        return path
353
354    def set_access_token(self, token: str) -> None:
355        self._access_token = token
356
357    def set_refresh_handler(self, handler: Callable[[], str]) -> None:
358        self._on_refresh_token = handler
359
360    def _do_fetch(
361        self,
362        path: str,
363        *,
364        method: str = "GET",
365        body: Any = None,
366        headers: dict[str, str] | None = None,
367        query: dict[str, Any] | None = None,
368    ) -> httpx.Response:
369        token = self._get_token()
370        url = f"{self._base_url}{self._transform_path(path)}"
371
372        req_headers = {
373            **self._default_headers,
374            "Content-Type": "application/json",
375        }
376        if token:
377            req_headers["Authorization"] = f"Bearer {token}"
378        if headers:
379            req_headers.update(headers)
380
381        params = None
382        if query:
383            params = {k: v for k, v in query.items() if v is not None}
384
385        return self._client.request(
386            method,
387            url,
388            json=body if body is not None and method not in ("GET", "HEAD") else None,
389            headers=req_headers,
390            params=params,
391        )
392
393    def _execute(
394        self,
395        path: str,
396        *,
397        method: str = "GET",
398        body: Any = None,
399        headers: dict[str, str] | None = None,
400        query: dict[str, Any] | None = None,
401    ) -> httpx.Response:
402        auth_prefix = f"{DEFAULT_API_PREFIX}/auth/"
403        if self._refresh_only and not path.startswith(auth_prefix):
404            raise RuntimeError(
405                f"Refresh-only HTTP client cannot make requests outside {auth_prefix}"
406            )
407
408        original_token = self._get_token()
409        response = self._do_fetch(path, method=method, body=body, headers=headers, query=query)
410
411        if (
412            response.status_code == 401
413            and self._on_refresh_token
414            and not path.startswith(auth_prefix)
415        ):
416            try:
417                with self._refresh_lock:
418                    if self._get_token() == original_token:
419                        self._access_token = self._on_refresh_token()
420            except Exception:
421                pass
422            else:
423                response = self._do_fetch(
424                    path, method=method, body=body, headers=headers, query=query
425                )
426
427        if response.status_code >= 400:
428            raw_data: dict[str, Any] = {}
429            try:
430                raw_data = response.json()
431            except Exception:
432                pass
433            error_code, message = _parse_error(raw_data, response.status_code)
434            raise ApiError(response.status_code, error_code, message, raw_data)
435
436        return response
437
438    @overload
439    def request(
440        self,
441        path: str,
442        *,
443        method: str = "GET",
444        body: Any = None,
445        headers: dict[str, str] | None = None,
446        query: dict[str, Any] | None = None,
447        response_type: type[T],
448    ) -> T: ...
449
450    @overload
451    def request(
452        self,
453        path: str,
454        *,
455        method: str = "GET",
456        body: Any = None,
457        headers: dict[str, str] | None = None,
458        query: dict[str, Any] | None = None,
459        response_type: None = None,
460    ) -> Any: ...
461
462    def request(
463        self,
464        path: str,
465        *,
466        method: str = "GET",
467        body: Any = None,
468        headers: dict[str, str] | None = None,
469        query: dict[str, Any] | None = None,
470        response_type: type[T] | None = None,
471    ) -> T | Any:
472        """Issue a request and return the JSON body.
473
474        With `response_type`, the body is validated into that type (a Pydantic
475        model or a generic alias like list[Model]); a bodyless 204 then raises
476        ValidationError, since the operation promised a typed body. Without
477        `response_type`, the raw parsed JSON is returned, or None on 204.
478        """
479        response = self._execute(path, method=method, body=body, headers=headers, query=query)
480
481        if response.status_code == 204:
482            if response_type is None:
483                return None
484            # A 204 on an operation that promises a typed body is a server
485            # contract violation; validating None fails loudly here instead
486            # of surfacing later as an AttributeError far from the call.
487            return _type_adapter(response_type).validate_python(None)
488
489        raw = response.json()
490        if response_type is None:
491            return raw
492        return _type_adapter(response_type).validate_python(raw)
493
494    def request_raw(
495        self,
496        path: str,
497        *,
498        method: str = "GET",
499        body: Any = None,
500        headers: dict[str, str] | None = None,
501        query: dict[str, Any] | None = None,
502    ) -> dict[str, Any]:
503        response = self._execute(path, method=method, body=body, headers=headers, query=query)
504
505        return {
506            "content": response.content,
507            "mime_type": response.headers.get("content-type", "text/plain"),
508        }
509
510    def stream_sse_sync(
511        self,
512        path: str,
513        *,
514        method: str = "GET",
515        body: Any = None,
516        headers: dict[str, str] | None = None,
517        query: dict[str, Any] | None = None,
518    ) -> Iterator[dict[str, Any]]:
519        """Synchronous counterpart to :meth:`HttpClient.stream_sse`.
520
521        Backs the generated sync ``stream()`` methods. Raises :class:`ApiError`
522        on a non-2xx response before the stream opens.
523        """
524        auth_prefix = f"{DEFAULT_API_PREFIX}/auth/"
525        if self._refresh_only and not path.startswith(auth_prefix):
526            raise RuntimeError(
527                f"Refresh-only HTTP client cannot make requests outside {auth_prefix}"
528            )
529
530        url = f"{self._base_url}{self._transform_path(path)}"
531        sends_body = body is not None and method not in ("GET", "HEAD")
532        req_headers = {**self._default_headers, "Accept": "text/event-stream"}
533        if sends_body:
534            req_headers["Content-Type"] = "application/json"
535        token = self._get_token()
536        if token:
537            req_headers["Authorization"] = f"Bearer {token}"
538        if headers:
539            req_headers.update(headers)
540        params = {k: v for k, v in (query or {}).items() if v is not None} or None
541
542        with self._client.stream(
543            method,
544            url,
545            json=body if sends_body else None,
546            headers=req_headers,
547            params=params,
548        ) as response:
549            if response.status_code >= 400:
550                response.read()
551                raw: dict[str, Any] = {}
552                try:
553                    raw = response.json()
554                except Exception:
555                    pass
556                code, message = _parse_error(raw, response.status_code)
557                raise ApiError(response.status_code, code, message, raw)
558
559            event: str | None = None
560            data_lines: list[str] = []
561            for line in response.iter_lines():
562                if line == "":
563                    parsed = _build_sse_event(event, data_lines)
564                    if parsed is not None:
565                        yield parsed
566                    event, data_lines = None, []
567                elif line.startswith("event:"):
568                    event = line[6:].strip()
569                elif line.startswith("data:"):
570                    data_lines.append(line[5:].strip())
571            parsed = _build_sse_event(event, data_lines)
572            if parsed is not None:
573                yield parsed
574
575    def close(self) -> None:
576        self._client.close()
577
578
579def _build_sse_event(event: str | None, data_lines: list[str]) -> dict[str, Any] | None:
580    """Assemble one SSE frame into ``{"event", "data"}``; ``None`` if empty."""
581    if event is None and not data_lines:
582        return None
583    raw = "\n".join(data_lines)
584    try:
585        data: Any = json.loads(raw)
586    except Exception:
587        data = raw
588    return {"event": event or "message", "data": data}
589
590
591def _parse_error(raw_data: dict[str, Any], status: int) -> tuple[str, str]:
592    error = raw_data.get("error")
593    if isinstance(error, dict):
594        code = error.get("code") or error.get("type") or "unknown_error"
595        message = error.get("message") or f"HTTP {status}"
596        return code, message
597    error_str = error if isinstance(error, str) else None
598    message = raw_data.get("message") or error_str or f"HTTP {status}"
599    return error_str or "unknown_error", message