Add with_retry helper for retrying Matrix API calls

matrix/api.py

@@ -1,21 +1,23 @@
 import asyncio
-from typing import Awaitable, Iterable, TypeVar
+from typing import Awaitable, Callable, Iterable, TypeVar
 
 from nio import ErrorResponse, Response
 
-from matrix.errors import MatrixError
+from matrix.errors import MatrixApiError
 
 T = TypeVar("T", bound=Response)
 R = TypeVar("R")
 
 
 async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T:
-    """Await `coro`, translating any failure into a `MatrixError`.
+    """Await `coro`, translating any failure into a `MatrixApiError`.
 
     matrix-nio's `AsyncClient` methods don't raise on API-level errors; they
     return an `ErrorResponse` instead of raising. This wraps a single call so
     both transport-level exceptions and nio `ErrorResponse` results become a
-    `MatrixError` carrying `error_message`.
+    `MatrixApiError` carrying `error_message`. When the `ErrorResponse` is a
+    rate limit with a `retry_after_ms`, that value is carried through
+    unchanged on `MatrixApiError.retry_after_ms`.
 
     ## Example
 
@@ -29,14 +31,54 @@ async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T:
     try:
         response = await coro
     except Exception as e:
-        raise MatrixError(f"{error_message}: {e}") from e
+        raise MatrixApiError(f"{error_message}: {e}") from e
 
     if isinstance(response, ErrorResponse):
-        raise MatrixError(f"{error_message}: {response}")
+        raise MatrixApiError(
+            f"{error_message}: {response}", retry_after_ms=response.retry_after_ms
+        )
 
     return response
 
 
+async def with_retry(
+    func: Callable[[], Awaitable[R]],
+    /,
+    *,
+    retries: int = 3,
+    base_delay: float = 1.0,
+) -> R:
+    """Call `func`, retrying on `MatrixApiError` with backoff.
+
+    Retries up to `retries` times after the initial attempt (so up to
+    `retries + 1` total calls). If the raised `MatrixApiError` carries a
+    `retry_after_ms` (from the Matrix server), that wait is used exactly
+    (converted to seconds); otherwise falls back to exponential delay
+    (`base_delay * 2 ** attempt`). Only `MatrixApiError` triggers a retry;
+    any other exception propagates.
+
+    ## Example
+
+    ```python
+    message = await with_retry(lambda: room.send(content, notice=notice))
+    ```
+    """
+    for attempt in range(retries + 1):
+        try:
+            return await func()
+        except MatrixApiError as e:
+            if attempt == retries:
+                raise
+
+            delay = base_delay * 2**attempt
+            if e.retry_after_ms is not None:
+                delay = e.retry_after_ms / 1000
+
+            await asyncio.sleep(delay)
+
+    raise AssertionError("unreachable")
+
+
 async def bounded_gather(
     coros: Iterable[Awaitable[R]],
     /,

matrix/errors.py

@@ -13,6 +13,12 @@ class MatrixError(Exception):
     pass
 
 
+class MatrixApiError(MatrixError):
+    def __init__(self, message: str, *, retry_after_ms: int | None = None):
+        super().__init__(message)
+        self.retry_after_ms = retry_after_ms
+
+
 class RoomNotFoundError(MatrixError):
     pass
 

tests/test_api.py

@@ -1,11 +1,12 @@
 import asyncio
 import inspect
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from nio import RoomSendResponse, RoomSendError
 
-from matrix.errors import MatrixError
-from matrix.api import bounded_gather, matrix_call
+from matrix.errors import ConfigError, MatrixApiError, MatrixError
+from matrix.api import bounded_gather, matrix_call, with_retry
 
 
 @pytest.mark.asyncio
@@ -19,11 +20,11 @@ async def call():
 
 
 @pytest.mark.asyncio
-async def test_matrix_call_with_transport_exception__expect_matrix_error():
+async def test_matrix_call_with_transport_exception__expect_matrix_api_error():
     async def call():
         raise Exception("Network error")
 
-    with pytest.raises(MatrixError, match="Failed to send message: Network error"):
+    with pytest.raises(MatrixApiError, match="Failed to send message: Network error"):
         await matrix_call(call(), error_message="Failed to send message")
 
 
@@ -36,6 +37,30 @@ async def call():
         await matrix_call(call(), error_message="Failed to send message")
 
 
+@pytest.mark.asyncio
+async def test_matrix_call_with_rate_limited_error_response__expect_retry_after_ms_set():
+    async def call():
+        return RoomSendError(
+            "too many requests", "M_LIMIT_EXCEEDED", retry_after_ms=5000
+        )
+
+    with pytest.raises(MatrixApiError) as exc_info:
+        await matrix_call(call(), error_message="Failed to send message")
+
+    assert exc_info.value.retry_after_ms == 5000
+
+
+@pytest.mark.asyncio
+async def test_matrix_call_with_error_response_without_retry_after__expect_retry_after_ms_none():
+    async def call():
+        return RoomSendError("not allowed", "M_FORBIDDEN")
+
+    with pytest.raises(MatrixApiError) as exc_info:
+        await matrix_call(call(), error_message="Failed to send message")
+
+    assert exc_info.value.retry_after_ms is None
+
+
 def test_matrix_call_requires_keyword_error_message__expect_type_error():
     with pytest.raises(TypeError):
         matrix_call(None, "Failed to send message")
@@ -46,6 +71,103 @@ def test_matrix_call_requires_positional_coro__expect_type_error():
         matrix_call(coro=None, error_message="Failed to send message")
 
 
+@pytest.mark.asyncio
+async def test_with_retry_with_success_first_try__expect_response_returned():
+    call = AsyncMock(return_value="ok")
+
+    result = await with_retry(lambda: call())
+
+    assert result == "ok"
+    assert call.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_with_retry_with_transient_matrix_error_then_success__expect_response_returned_after_retry():
+    call = AsyncMock(side_effect=[MatrixApiError("rate limited"), "ok"])
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep:
+        result = await with_retry(lambda: call())
+
+    assert result == "ok"
+    assert call.await_count == 2
+    sleep.assert_awaited_once_with(1.0)
+
+
+@pytest.mark.asyncio
+async def test_with_retry_exhausts_retries__expect_matrix_api_error_raised():
+    call = AsyncMock(side_effect=MatrixApiError("still failing"))
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock):
+        with pytest.raises(MatrixApiError, match="still failing"):
+            await with_retry(lambda: call(), retries=2)
+
+    assert call.await_count == 3
+
+
+@pytest.mark.asyncio
+async def test_with_retry_with_non_matrix_exception__expect_immediate_raise_no_retry():
+    call = AsyncMock(side_effect=Exception("boom"))
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep:
+        with pytest.raises(Exception, match="boom"):
+            await with_retry(lambda: call())
+
+    assert call.await_count == 1
+    sleep.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_with_retry_with_non_api_matrix_error__expect_immediate_raise_no_retry():
+    call = AsyncMock(side_effect=ConfigError("token"))
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep:
+        with pytest.raises(ConfigError):
+            await with_retry(lambda: call())
+
+    assert call.await_count == 1
+    sleep.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_with_retry_uses_exponential_backoff__expect_increasing_delays():
+    call = AsyncMock(
+        side_effect=[
+            MatrixApiError("fail 1"),
+            MatrixApiError("fail 2"),
+            MatrixApiError("fail 3"),
+            "ok",
+        ]
+    )
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep:
+        result = await with_retry(lambda: call(), retries=3, base_delay=1.0)
+
+    assert result == "ok"
+    assert [call_args.args[0] for call_args in sleep.await_args_list] == [
+        1.0,
+        2.0,
+        4.0,
+    ]
+
+
+@pytest.mark.asyncio
+async def test_with_retry_with_retry_after_ms__expect_server_delay_honored():
+    call = AsyncMock(
+        side_effect=[MatrixApiError("rate limited", retry_after_ms=2500), "ok"]
+    )
+
+    with patch("matrix.api.asyncio.sleep", new_callable=AsyncMock) as sleep:
+        result = await with_retry(lambda: call())
+
+    assert result == "ok"
+    sleep.assert_awaited_once_with(2.5)
+
+
+def test_with_retry_requires_positional_func__expect_type_error():
+    with pytest.raises(TypeError):
+        with_retry(func=lambda: None)
+
+
 @pytest.mark.asyncio
 async def test_bounded_gather_with_success__expect_results_in_order():
     async def value(n):