improve API client

pyproject.toml

@@ -51,3 +51,4 @@ select = [
 [tool.pyright]
 include = ["src"]
 exclude = ["**/__pycache__"]
+strict = ["src/fg.py"]

src/app.py

@@ -1,5 +1,6 @@
 import dataclasses as dc
 import io
+from functools import cache
 from typing import Any
 
 import gradio as gr
@@ -24,6 +25,7 @@ with env.prefixed("ERASER_"):
     CA_BUNDLE: str | None = env.str("CA_BUNDLE", None)
 
 
+@cache
 def _ctx() -> EditorAPIContext:
     assert API_USER is not None
     assert API_PASSWORD is not None
@@ -51,13 +53,7 @@ class ProcessParams:
 async def _process(ctx: EditorAPIContext, params: ProcessParams) -> Image.Image:
     with io.BytesIO() as f:
         params.image.save(f, format="JPEG")
-        async with ctx as client:
-            response = await client.post(
-                f"{ctx.uri}/state/upload",
-                files={"file": f},
-                headers=ctx.auth_headers,
-            )
-    response.raise_for_status()
+        response = await ctx.request("POST", "state/upload", files={"file": f})
     st_input = response.json()["state"]
 
     if params.bbox:
@@ -74,13 +70,11 @@ async def _process(ctx: EditorAPIContext, params: ProcessParams) -> Image.Image:
     st_mask = await ctx.call_skill(f"segment/{segment_input_st}", segment_params)
     st_erased = await ctx.call_skill(f"erase/{st_input}/{st_mask}", {"mode": "free"})
 
-    async with ctx as client:
-        response = await client.get(
-            f"{ctx.uri}/state/image/{st_erased}",
-            params={"format": "JPEG", "resolution": "DISPLAY"},
-            headers=ctx.auth_headers,
-        )
-    response.raise_for_status()
+    response = await ctx.request(
+        "GET",
+        f"state/image/{st_erased}",
+        params={"format": "JPEG", "resolution": "DISPLAY"},
+    )
     f = io.BytesIO()
     f.write(response.content)
     f.seek(0)

src/fg.py

@@ -1,18 +1,46 @@
 import asyncio
 import dataclasses as dc
 import json
+import logging
 from collections import defaultdict
-from collections.abc import Awaitable, Callable
-from typing import Any, Literal
+from collections.abc import Awaitable, Callable, Mapping
+from typing import Any, Literal, cast
 
 import httpx
 import httpx_sse
+from httpx._types import QueryParamTypes, RequestFiles
+
+logger = logging.getLogger(__name__)
 
 Priority = Literal["low", "standard", "high"]
 
 
-def _new_future() -> asyncio.Future[Any]:
-    return asyncio.get_running_loop().create_future()
+class SSELoopStopped(RuntimeError):
+    pass
+
+
+class Futures[T]:
+    @classmethod
+    def create_future(cls) -> asyncio.Future[T]:
+        return asyncio.get_running_loop().create_future()
+
+    def __init__(self, capacity: int = 256) -> None:
+        self.futures = defaultdict[str, asyncio.Future[T]](self.create_future)
+        self.capacity = capacity
+
+    def cull(self) -> None:
+        while len(self.futures) >= self.capacity:
+            del self.futures[next(iter(self.futures))]
+
+    def __getitem__(self, key: str) -> asyncio.Future[T]:
+        self.cull()
+        return self.futures[key]
+
+    def __delitem__(self, key: str) -> None:
+        try:
+            del self.futures[key]
+        except KeyError:
+            pass
 
 
 @dc.dataclass(kw_only=True)
@@ -23,18 +51,33 @@ class EditorAPIContext:
     priority: Priority = "standard"
     token: str | None = None
     verify: bool | str = True
-    _client: httpx.AsyncClient | None = None
+    default_timeout: float = 60.0
+    logger: logging.Logger = logger
+    max_sse_failures: int = 5
 
-    sse_futures: dict[str, asyncio.Future[dict[str, Any]]] = dc.field(default_factory=lambda: defaultdict(_new_future))
+    _client: httpx.AsyncClient | None = None
+    _client_ctx_depth: int = 0
+    _sse_futures: Futures[dict[str, Any]] = dc.field(default_factory=Futures)
+    _sse_task: asyncio.Task[None] | None = None
+    _sse_failures: int = 0
+    _sse_last_event_id: str = ""
+    _sse_retry_ms: int = 0
 
     async def __aenter__(self) -> httpx.AsyncClient:
         if self._client:
+            assert self._client_ctx_depth > 0
+            self._client_ctx_depth += 1
             return self._client
+        assert self._client_ctx_depth == 0
         self._client = httpx.AsyncClient(verify=self.verify)
+        self._client_ctx_depth = 1
         return self._client
 
     async def __aexit__(self, *args: Any) -> None:
-        if self._client:
+        if (not self._client) or self._client_ctx_depth <= 0:
+            raise RuntimeError("unbalanced __aexit__")
+        self._client_ctx_depth -= 1
+        if self._client_ctx_depth == 0:
             await self._client.__aexit__(*args)
             self._client = None
 
@@ -49,62 +92,153 @@ class EditorAPIContext:
                 f"{self.uri}/auth/login",
                 json={"username": self.user, "password": self.password},
             )
-            response.raise_for_status()
-            self.token = response.json()["token"]
+        response.raise_for_status()
+        self.logger.debug(f"logged in as {self.user}")
+        self.token = response.json()["token"]
+
+    async def request(
+        self,
+        method: Literal["GET", "POST"],
+        url: str,
+        files: RequestFiles | None = None,
+        params: QueryParamTypes | None = None,
+        json: dict[str, Any] | None = None,
+        headers: Mapping[str, str] | None = None,
+        raise_for_status: bool = True,
+    ) -> httpx.Response:
+        async def _q() -> httpx.Response:
+            return await client.request(
+                method,
+                f"{self.uri}/{url}",
+                headers=dict(headers or {}) | self.auth_headers,
+                files=files,
+                params=params,
+                json=json,
+            )
 
-    async def sse_loop(self) -> None:
         async with self as client:
-            response = await client.post(f"{self.uri}/sub-auth", headers=self.auth_headers)
-        response.raise_for_status()
+            r = await _q()
+            if r.status_code == 401:
+                self.logger.debug("renewing token")
+                await self.login()
+                r = await _q()
+
+        if raise_for_status:
+            r.raise_for_status()
+        return r
+
+    @classmethod
+    def decode_json(cls, data: str) -> dict[str, Any] | None:
+        try:
+            r = json.loads(data)
+        except json.JSONDecodeError:
+            return None
+        if type(r) is not dict:
+            return None
+        return cast(dict[str, Any], r)
+
+    async def _sse_loop(self) -> None:
+        response = await self.request("POST", "sub-auth")
         sub_token = response.json()["token"]
         url = f"{self.uri}/sub/{sub_token}"
+        headers = {"Accept": "text/event-stream"}
+        if self._sse_last_event_id:
+            retry_ms = self._sse_retry_ms + 1000 * 2**self._sse_failures
+            self.logger.info(f"resuming SSE from event {self._sse_last_event_id} in {retry_ms} ms")
+            await asyncio.sleep(retry_ms / 1000)
+            headers["Last-Event-ID"] = self._sse_last_event_id
         async with (
             httpx.AsyncClient(timeout=None, verify=self.verify) as c,
-            httpx_sse.aconnect_sse(c, "GET", url) as es,
+            httpx_sse.aconnect_sse(c, "GET", url, headers=headers) as es,
         ):
-            future = self.sse_futures["_sse_loop"]
-            future.set_result({"status": "ok"})
-            async for sse in es.aiter_sse():
-                jdata = json.loads(sse.data)
-                future = self.sse_futures[jdata["state"]]
-                future.set_result(jdata)
-
-    async def sse_await(self, state_id: str, timeout: float = 60.0) -> None:
-        future = self.sse_futures[state_id]
-        jdata = await asyncio.wait_for(future, timeout=timeout)
-        if jdata["status"] != "ok":
-            print("ERROR", jdata)
-        assert jdata["status"] == "ok"
-        del self.sse_futures[state_id]
+            es.response.raise_for_status()
+            self._sse_futures["_sse_loop"].set_result({"status": "ok"})
+            try:
+                async for sse in es.aiter_sse():
+                    self._sse_last_event_id = sse.id
+                    self._sse_retry_ms = sse.retry or 0
+                    jdata = self.decode_json(sse.data)
+                    if (jdata is None) or ("state" not in jdata):
+                        # Note: when the server restarts we typically get an
+                        # empty string here, then the loop exits.
+                        self.logger.warning(f"unexpected SSE data: {sse.data}")
+                        continue
+                    self._sse_futures[jdata["state"]].set_result(jdata)
+            except asyncio.CancelledError:
+                pass
 
-    async def get_meta(self, state_id: str) -> dict[str, Any]:
-        async with self as client:
-            response = await client.get(
-                f"{self.uri}/state/meta/{state_id}",
-                headers=self.auth_headers,
+    async def sse_start(self) -> None:
+        assert self._sse_task is None
+        self._sse_last_event_id = ""
+        self._sse_retry_ms = 0
+        self._sse_task = asyncio.create_task(self._sse_loop())
+        await self.sse_await("_sse_loop")
+        self._sse_failures = 0
+
+    async def sse_recover(self) -> bool:
+        while True:
+            if self._sse_failures > self.max_sse_failures:
+                return False
+            self._sse_task = asyncio.create_task(self._sse_loop())
+            try:
+                await self.sse_await("_sse_loop")
+                return True
+            except SSELoopStopped:
+                pass
+
+    async def sse_stop(self) -> None:
+        assert self._sse_task
+        self._sse_task.cancel()
+        await self._sse_task
+        self._sse_task = None
+
+    async def sse_await(self, state_id: str, timeout: float | None = None) -> None:
+        assert self._sse_task
+        future = self._sse_futures[state_id]
+
+        while True:
+            done, _ = await asyncio.wait(
+                {future, self._sse_task},
+                timeout=timeout or self.default_timeout,
+                return_when=asyncio.FIRST_COMPLETED,
             )
-            response.raise_for_status()
-            return response.json()
+            if not done:
+                raise TimeoutError(f"state {state_id} timed out after {timeout}")
+            if self._sse_task in done:
+                self._sse_failures += 1
+                if state_id != "_sse_loop" and (await self.sse_recover()):
+                    self._sse_failures = 0
+                    continue
+                exception = self._sse_task.exception()
+                raise SSELoopStopped(f"SSE loop stopped while waiting for state {state_id}") from exception
+            break
 
-    async def run_one[Tin, Tout](
+        assert done == {future}
+
+        jdata = future.result()
+        del self._sse_futures[state_id]
+        assert jdata["status"] == "ok", f"state {state_id} is {jdata['status']}"
+
+    async def get_meta(self, state_id: str) -> dict[str, Any]:
+        response = await self.request("GET", f"state/meta/{state_id}")
+        return response.json()
+
+    async def _run_one[Tin, Tout](
         self,
         co: Callable[["EditorAPIContext", Tin], Awaitable[Tout]],
         params: Tin,
     ) -> Tout:
-        await self.login()
-        async with asyncio.TaskGroup() as tg:
-            sse_task = tg.create_task(self.sse_loop())
-
-            async def outer_co(params: Tin) -> Tout:
-                # _sse_loop is a fake event to wait until the SSE loop is properly setup.
-                await self.sse_await("_sse_loop")
-                r = await co(self, params)
-                sse_task.cancel()
-                return r
-
-            r = tg.create_task(outer_co(params))
-
-        return r.result()
+        # This wraps the coroutine in the SSE loop.
+        # This is mostly useful if you use synchronous Python,
+        # otherwise you can call the functions directly.
+        if not self.token:
+            await self.login()
+        await self.sse_start()
+        try:
+            r = await co(self, params)
+            return r
+        finally:
+            await self.sse_stop()
 
     def run_one_sync[Tin, Tout](
         self,
@@ -116,18 +250,16 @@ class EditorAPIContext:
         except RuntimeError:
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
+        return loop.run_until_complete(self._run_one(co, params))
 
-        return loop.run_until_complete(self.run_one(co, params))
-
-    async def call_skill(self, uri: str, params: dict[str, Any] | None) -> str:
+    async def call_skill(
+        self,
+        uri: str,
+        params: dict[str, Any] | None,
+        timeout: float | None = None,
+    ) -> str:
         params = {"priority": self.priority} | (params or {})
-        async with self as client:
-            response = await client.post(
-                f"{self.uri}/skills/{uri}",
-                json=params,
-                headers=self.auth_headers,
-            )
-        response.raise_for_status()
+        response = await self.request("POST", f"skills/{uri}", json=params)
         state_id = response.json()["state"]
-        await self.sse_await(state_id)
+        await self.sse_await(state_id, timeout=timeout)
         return state_id

typings/pillow_heif/__init__.pyi

@@ -0,0 +1,2 @@
+def register_heif_opener() -> None: ...
+def register_avif_opener() -> None: ...