Retry streamed audit log transfers

Author: yummybomb

src/kernel/lib/audit_log_download.py

@@ -3,7 +3,7 @@
 import time
 import hashlib
 import inspect
-from typing import BinaryIO, Callable, Optional, Protocol, Awaitable
+from typing import BinaryIO, Callable, Optional, Protocol, Awaitable, ContextManager, AsyncContextManager
 from dataclasses import dataclass
 
 import anyio
@@ -51,8 +51,8 @@ async def read(self) -> bytes: ...
     async def close(self) -> None: ...
 
 
-SyncFetchChunk = Callable[[Optional[str]], _SyncChunkResponse]
-AsyncFetchChunk = Callable[[Optional[str]], Awaitable[_AsyncChunkResponse]]
+SyncFetchChunk = Callable[[Optional[str]], ContextManager[_SyncChunkResponse]]
+AsyncFetchChunk = Callable[[Optional[str]], AsyncContextManager[_AsyncChunkResponse]]
 ProgressCallback = Callable[[AuditLogDownloadProgress], None]
 AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]]
 
@@ -147,39 +147,41 @@ def _fetch_verified_chunk(
     fetch_chunk: SyncFetchChunk, cursor: str | None, max_transfer_retries: int
 ) -> tuple[bytes, httpx.Headers]:
     for retries in range(max_transfer_retries + 1):
-        response = fetch_chunk(cursor)
-        try:
+        transfer_error: Exception | None = None
+        with fetch_chunk(cursor) as response:
             try:
                 body = response.read()
                 headers = httpx.Headers(response.headers)
-            finally:
-                response.close()
-            _verify_checksum(body, headers)
-            return body, headers
-        except Exception:
-            if retries == max_transfer_retries:
-                raise
-            time.sleep(min(2**retries, _MAX_RETRY_DELAY))
+                _verify_checksum(body, headers)
+            except Exception as error:
+                transfer_error = error
+            else:
+                return body, headers
+        assert transfer_error is not None
+        if retries == max_transfer_retries:
+            raise transfer_error
+        time.sleep(min(2**retries, _MAX_RETRY_DELAY))
     raise AssertionError("unreachable")
 
 
 async def _async_fetch_verified_chunk(
     fetch_chunk: AsyncFetchChunk, cursor: str | None, max_transfer_retries: int
 ) -> tuple[bytes, httpx.Headers]:
     for retries in range(max_transfer_retries + 1):
-        response = await fetch_chunk(cursor)
-        try:
+        transfer_error: Exception | None = None
+        async with fetch_chunk(cursor) as response:
             try:
                 body = await response.read()
                 headers = httpx.Headers(response.headers)
-            finally:
-                await response.close()
-            await to_thread.run_sync(_verify_checksum, body, headers)
-            return body, headers
-        except Exception:
-            if retries == max_transfer_retries:
-                raise
-            await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY))
+                await to_thread.run_sync(_verify_checksum, body, headers)
+            except Exception as error:
+                transfer_error = error
+            else:
+                return body, headers
+        assert transfer_error is not None
+        if retries == max_transfer_retries:
+            raise transfer_error
+        await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY))
     raise AssertionError("unreachable")
 
 
@@ -199,7 +201,12 @@ def _parse_chunk_headers(headers: httpx.Headers, current_cursor: str | None) ->
     has_more = has_more_value == "true"
 
     row_count = headers.get("x-row-count")
-    if row_count is None or not row_count.isascii() or not row_count.isdecimal():
+    if (
+        row_count is None
+        or not row_count.isascii()
+        or not row_count.isdecimal()
+        or len(row_count) > len(str(_MAX_CHUNK_ROWS))
+    ):
         raise AuditLogDownloadError("response missing or invalid X-Row-Count header")
     rows = int(row_count)
     if rows > _MAX_CHUNK_ROWS:

src/kernel/resources/audit_logs.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Union, BinaryIO
+from typing import Union, BinaryIO, ContextManager, AsyncContextManager
 from datetime import datetime
 from typing_extensions import Literal
 
@@ -258,8 +258,8 @@ def download(
         when the completed export must be published atomically.
         """
 
-        def fetch_chunk(cursor: str | None) -> BinaryAPIResponse:
-            return self.export_chunk(
+        def fetch_chunk(cursor: str | None) -> ContextManager[StreamedBinaryAPIResponse]:
+            return self.with_streaming_response.export_chunk(
                 end=end,
                 start=start,
                 auth_strategy=auth_strategy,
@@ -502,8 +502,8 @@ async def download(
         when the completed export must be published atomically.
         """
 
-        async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse:
-            return await self.export_chunk(
+        def fetch_chunk(cursor: str | None) -> AsyncContextManager[AsyncStreamedBinaryAPIResponse]:
+            return self.with_streaming_response.export_chunk(
                 end=end,
                 start=start,
                 auth_strategy=auth_strategy,

tests/test_audit_log_download.py

@@ -3,7 +3,8 @@
 import hashlib
 import threading
 from io import BytesIO
-from typing import BinaryIO, cast
+from typing import BinaryIO, Iterator, AsyncIterator, cast
+from typing_extensions import override
 
 import httpx
 import pytest
@@ -13,15 +14,37 @@
 from kernel.lib import audit_log_download
 
 
-def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response:
+def chunk_headers(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> dict[str, str]:
     headers = {
         "x-content-sha256": hashlib.sha256(body).hexdigest(),
         "x-has-more": str(has_more).lower(),
         "x-row-count": str(rows),
     }
     if next_cursor is not None:
         headers["x-next-cursor"] = next_cursor
-    return httpx.Response(200, content=body, headers=headers)
+    return headers
+
+
+def chunk_response(body: bytes, *, rows: int, has_more: bool, next_cursor: str | None = None) -> httpx.Response:
+    return httpx.Response(
+        200,
+        content=body,
+        headers=chunk_headers(body, rows=rows, has_more=has_more, next_cursor=next_cursor),
+    )
+
+
+class BrokenSyncByteStream(httpx.SyncByteStream):
+    @override
+    def __iter__(self) -> Iterator[bytes]:
+        yield b"partial"
+        raise httpx.ReadError("truncated response body")
+
+
+class BrokenAsyncByteStream(httpx.AsyncByteStream):
+    @override
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        yield b"partial"
+        raise httpx.ReadError("truncated response body")
 
 
 def test_download_writes_verified_chunks(client: Kernel, respx_mock: MockRouter) -> None:
@@ -98,6 +121,33 @@ async def on_progress(_: object) -> None:
     assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids)
 
 
+async def test_async_download_retries_body_read_failure(
+    async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    async def no_sleep(_: float) -> None:
+        pass
+
+    monkeypatch.setattr(audit_log_download.anyio, "sleep", no_sleep)
+    headers = chunk_headers(b"good", rows=1, has_more=False)
+    route = respx_mock.get("/audit-logs/export/chunk").mock(
+        side_effect=[
+            httpx.Response(200, headers=headers, stream=BrokenAsyncByteStream()),
+            chunk_response(b"good", rows=1, has_more=False),
+        ]
+    )
+    destination = BytesIO()
+
+    await async_client.with_options(max_retries=0).audit_logs.download(
+        to=destination,
+        start="2026-06-01T00:00:00Z",
+        end="2026-06-02T00:00:00Z",
+        max_transfer_retries=1,
+    )
+
+    assert destination.getvalue() == b"good"
+    assert route.call_count == 2
+
+
 async def test_async_download_retries_checksum_mismatch(
     async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -154,7 +204,7 @@ async def test_async_download_respects_disabled_http_retries(async_client: Async
     assert route.call_count == 1
 
 
-@pytest.mark.parametrize("row_count", ["", "1.0", "50001"])
+@pytest.mark.parametrize("row_count", ["", "1.0", "50001", "9" * 5000])
 def test_download_rejects_invalid_row_count(row_count: str, client: Kernel, respx_mock: MockRouter) -> None:
     response = chunk_response(b"chunk", rows=1, has_more=False)
     response.headers["x-row-count"] = row_count
@@ -202,6 +252,33 @@ def test_download_rejects_invalid_transfer_retry_count(max_transfer_retries: obj
         )
 
 
+def test_download_retries_body_read_failure(
+    client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    def no_sleep(_: float) -> None:
+        pass
+
+    monkeypatch.setattr(audit_log_download.time, "sleep", no_sleep)
+    headers = chunk_headers(b"good", rows=1, has_more=False)
+    route = respx_mock.get("/audit-logs/export/chunk").mock(
+        side_effect=[
+            httpx.Response(200, headers=headers, stream=BrokenSyncByteStream()),
+            chunk_response(b"good", rows=1, has_more=False),
+        ]
+    )
+    destination = BytesIO()
+
+    client.with_options(max_retries=0).audit_logs.download(
+        to=destination,
+        start="2026-06-01T00:00:00Z",
+        end="2026-06-02T00:00:00Z",
+        max_transfer_retries=1,
+    )
+
+    assert destination.getvalue() == b"good"
+    assert route.call_count == 2
+
+
 def test_download_retries_checksum_mismatch(
     client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
 ) -> None: