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