Harden audit log downloads

Author: yummybomb

src/kernel/lib/audit_log_download.py

@@ -2,15 +2,17 @@
 
 import time
 import hashlib
+import inspect
 from typing import BinaryIO, Callable, Optional, Protocol, Awaitable
 from dataclasses import dataclass
 
 import anyio
 import httpx
+from anyio import to_thread
 
-from .._exceptions import KernelError, APIStatusError
+from .._exceptions import KernelError
 
-_DOWNLOAD_ATTEMPTS = 7
+_DEFAULT_MAX_TRANSFER_RETRIES = 6
 _MAX_RETRY_DELAY = 8.0
 
 
@@ -51,22 +53,30 @@ async def close(self) -> None: ...
 SyncFetchChunk = Callable[[Optional[str]], _SyncChunkResponse]
 AsyncFetchChunk = Callable[[Optional[str]], Awaitable[_AsyncChunkResponse]]
 ProgressCallback = Callable[[AuditLogDownloadProgress], None]
+AsyncProgressCallback = Callable[[AuditLogDownloadProgress], Optional[Awaitable[None]]]
 
 
 def download_audit_logs(
     fetch_chunk: SyncFetchChunk,
     destination: BinaryIO,
     *,
     on_progress: ProgressCallback | None = None,
+    max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES,
 ) -> AuditLogDownloadResult:
     if not callable(getattr(destination, "write", None)):
         raise TypeError("audit log download destination must provide write()")
+    _validate_max_transfer_retries(max_transfer_retries)
 
     cursor: str | None = None
     result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0)
+    seen_cursors: set[str] = set()
     while True:
-        body, headers = _fetch_verified_chunk(fetch_chunk, cursor)
+        body, headers = _fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries)
         chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor)
+        if has_more and next_cursor is not None:
+            if next_cursor in seen_cursors:
+                raise AuditLogDownloadError("response repeated X-Next-Cursor header")
+            seen_cursors.add(next_cursor)
         _write_chunk(destination, body)
 
         cursor = next_cursor
@@ -92,17 +102,24 @@ async def async_download_audit_logs(
     fetch_chunk: AsyncFetchChunk,
     destination: BinaryIO,
     *,
-    on_progress: ProgressCallback | None = None,
+    on_progress: AsyncProgressCallback | None = None,
+    max_transfer_retries: int = _DEFAULT_MAX_TRANSFER_RETRIES,
 ) -> AuditLogDownloadResult:
     if not callable(getattr(destination, "write", None)):
         raise TypeError("audit log download destination must provide write()")
+    _validate_max_transfer_retries(max_transfer_retries)
 
     cursor: str | None = None
     result = AuditLogDownloadResult(bytes_written=0, chunks=0, rows=0)
+    seen_cursors: set[str] = set()
     while True:
-        body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor)
+        body, headers = await _async_fetch_verified_chunk(fetch_chunk, cursor, max_transfer_retries)
         chunk_rows, next_cursor, has_more = _parse_chunk_headers(headers, cursor)
-        _write_chunk(destination, body)
+        if has_more and next_cursor is not None:
+            if next_cursor in seen_cursors:
+                raise AuditLogDownloadError("response repeated X-Next-Cursor header")
+            seen_cursors.add(next_cursor)
+        await to_thread.run_sync(_write_chunk, destination, body)
 
         cursor = next_cursor
         result = AuditLogDownloadResult(
@@ -111,51 +128,57 @@ async def async_download_audit_logs(
             rows=result.rows + chunk_rows,
         )
         if on_progress is not None:
-            on_progress(
+            callback_result = on_progress(
                 AuditLogDownloadProgress(
                     bytes_written=result.bytes_written,
                     chunks=result.chunks,
                     rows=result.rows,
                     chunk_rows=chunk_rows,
                 )
             )
+            if inspect.isawaitable(callback_result):
+                await callback_result
         if not has_more:
             return result
 
 
-def _fetch_verified_chunk(fetch_chunk: SyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]:
-    for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1):
+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:
-            response = fetch_chunk(cursor)
             try:
                 body = response.read()
                 headers = httpx.Headers(response.headers)
             finally:
                 response.close()
             _verify_checksum(body, headers)
             return body, headers
-        except Exception as error:
-            if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error):
+        except Exception:
+            if retries == max_transfer_retries:
                 raise
-            time.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY))
+            time.sleep(min(2**retries, _MAX_RETRY_DELAY))
     raise AssertionError("unreachable")
 
 
-async def _async_fetch_verified_chunk(fetch_chunk: AsyncFetchChunk, cursor: str | None) -> tuple[bytes, httpx.Headers]:
-    for attempt in range(1, _DOWNLOAD_ATTEMPTS + 1):
+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:
-            response = await fetch_chunk(cursor)
             try:
                 body = await response.read()
                 headers = httpx.Headers(response.headers)
             finally:
                 await response.close()
-            _verify_checksum(body, headers)
+            await to_thread.run_sync(_verify_checksum, body, headers)
             return body, headers
-        except Exception as error:
-            if attempt == _DOWNLOAD_ATTEMPTS or not _is_retryable(error):
+        except Exception:
+            if retries == max_transfer_retries:
                 raise
-            await anyio.sleep(min(2 ** (attempt - 1), _MAX_RETRY_DELAY))
+            await anyio.sleep(min(2**retries, _MAX_RETRY_DELAY))
     raise AssertionError("unreachable")
 
 
@@ -199,7 +222,6 @@ def _write_chunk(destination: BinaryIO, body: bytes) -> None:
         remaining = remaining[written:]
 
 
-def _is_retryable(error: Exception) -> bool:
-    if isinstance(error, APIStatusError):
-        return error.status_code == 429 or error.status_code >= 500
-    return True
+def _validate_max_transfer_retries(max_transfer_retries: int) -> None:
+    if max_transfer_retries < 0:
+        raise ValueError("max_transfer_retries must be non-negative")

src/kernel/resources/audit_logs.py

@@ -32,6 +32,7 @@
 from ..types.audit_log_entry import AuditLogEntry
 from ..lib.audit_log_download import (
     ProgressCallback,
+    AsyncProgressCallback,
     AuditLogDownloadResult,
     download_audit_logs,
     async_download_audit_logs,
@@ -243,6 +244,7 @@ def download(
         search_user_id: SequenceNotStr[str] | Omit = omit,
         service: str | Omit = omit,
         on_progress: ProgressCallback | None = None,
+        max_transfer_retries: int = 6,
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
@@ -251,13 +253,13 @@ def download(
         """Download a complete audit log export to a writable binary destination.
 
         The SDK verifies every chunk and retries transient transfer failures. It
-        does not close the destination.
+        does not close the destination. If the download fails, the destination
+        may contain a partial export; use a temporary file and atomic rename
+        when the completed export must be published atomically.
         """
 
-        resource = self._client.with_options(max_retries=0).audit_logs
-
         def fetch_chunk(cursor: str | None) -> BinaryAPIResponse:
-            return resource.export_chunk(
+            return self.export_chunk(
                 end=end,
                 start=start,
                 auth_strategy=auth_strategy,
@@ -275,7 +277,12 @@ def fetch_chunk(cursor: str | None) -> BinaryAPIResponse:
                 timeout=timeout,
             )
 
-        return download_audit_logs(fetch_chunk, to, on_progress=on_progress)
+        return download_audit_logs(
+            fetch_chunk,
+            to,
+            on_progress=on_progress,
+            max_transfer_retries=max_transfer_retries,
+        )
 
 
 class AsyncAuditLogsResource(AsyncAPIResource):
@@ -480,7 +487,8 @@ async def download(
         search: str | Omit = omit,
         search_user_id: SequenceNotStr[str] | Omit = omit,
         service: str | Omit = omit,
-        on_progress: ProgressCallback | None = None,
+        on_progress: AsyncProgressCallback | None = None,
+        max_transfer_retries: int = 6,
         extra_headers: Headers | None = None,
         extra_query: Query | None = None,
         extra_body: Body | None = None,
@@ -489,13 +497,13 @@ async def download(
         """Download a complete audit log export to a writable binary destination.
 
         The SDK verifies every chunk and retries transient transfer failures. It
-        does not close the destination.
+        does not close the destination. If the download fails, the destination
+        may contain a partial export; use a temporary file and atomic rename
+        when the completed export must be published atomically.
         """
 
-        resource = self._client.with_options(max_retries=0).audit_logs
-
         async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse:
-            return await resource.export_chunk(
+            return await self.export_chunk(
                 end=end,
                 start=start,
                 auth_strategy=auth_strategy,
@@ -513,7 +521,12 @@ async def fetch_chunk(cursor: str | None) -> AsyncBinaryAPIResponse:
                 timeout=timeout,
             )
 
-        return await async_download_audit_logs(fetch_chunk, to, on_progress=on_progress)
+        return await async_download_audit_logs(
+            fetch_chunk,
+            to,
+            on_progress=on_progress,
+            max_transfer_retries=max_transfer_retries,
+        )
 
 
 class AuditLogsResourceWithRawResponse:

tests/test_audit_log_download.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import hashlib
+import threading
 from io import BytesIO
+from typing import BinaryIO
 
 import httpx
 import pytest
 from respx import MockRouter
 
-from kernel import Kernel, AsyncKernel, AuditLogDownloadError
+from kernel import Kernel, AsyncKernel, InternalServerError, AuditLogDownloadError
 from kernel.lib import audit_log_download
 
 
@@ -57,25 +59,43 @@ def respond(request: httpx.Request) -> httpx.Response:
     assert cursors == [None, "next"]
 
 
-async def test_async_download_writes_verified_chunks(async_client: AsyncKernel, respx_mock: MockRouter) -> None:
+async def test_async_download_writes_verified_chunks(
+    async_client: AsyncKernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
     respx_mock.get("/audit-logs/export/chunk").mock(
         side_effect=[
             chunk_response(b"first", rows=2, has_more=True, next_cursor="next"),
             chunk_response(b"second", rows=1, has_more=False),
         ]
     )
+    thread_ids: list[int] = []
+    write_chunk = audit_log_download._write_chunk
+
+    def record_write(destination: BinaryIO, body: bytes) -> None:
+        thread_ids.append(threading.get_ident())
+        write_chunk(destination, body)
+
+    monkeypatch.setattr(audit_log_download, "_write_chunk", record_write)
     destination = BytesIO()
+    progress_called = False
+
+    async def on_progress(_: object) -> None:
+        nonlocal progress_called
+        progress_called = True
 
     result = await async_client.audit_logs.download(
         to=destination,
         start="2026-06-01T00:00:00Z",
         end="2026-06-02T00:00:00Z",
+        on_progress=on_progress,
     )
 
     assert destination.getvalue() == b"firstsecond"
     assert result.bytes_written == 11
     assert result.chunks == 2
     assert result.rows == 3
+    assert progress_called
+    assert thread_ids and all(thread_id != threading.get_ident() for thread_id in thread_ids)
 
 
 def test_download_retries_checksum_mismatch(
@@ -102,9 +122,7 @@ def no_sleep(_: float) -> None:
     assert route.call_count == 2
 
 
-def test_download_owns_http_retries(client: Kernel, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None:
-    delays: list[float] = []
-    monkeypatch.setattr(audit_log_download.time, "sleep", delays.append)
+def test_download_uses_client_http_retries(client: Kernel, respx_mock: MockRouter) -> None:
     route = respx_mock.get("/audit-logs/export/chunk").mock(
         side_effect=[
             httpx.Response(500, json={"message": "temporary failure"}),
@@ -119,7 +137,57 @@ def test_download_owns_http_retries(client: Kernel, respx_mock: MockRouter, monk
     )
 
     assert route.call_count == 2
-    assert delays == [1]
+
+
+def test_download_respects_disabled_http_retries(client: Kernel, respx_mock: MockRouter) -> None:
+    route = respx_mock.get("/audit-logs/export/chunk").mock(
+        return_value=httpx.Response(500, json={"message": "temporary failure"})
+    )
+
+    with pytest.raises(InternalServerError, match="500"):
+        client.with_options(max_retries=0).audit_logs.download(
+            to=BytesIO(),
+            start="2026-06-01T00:00:00Z",
+            end="2026-06-02T00:00:00Z",
+        )
+
+    assert route.call_count == 1
+
+
+def test_download_respects_transfer_retry_limit(client: Kernel, respx_mock: MockRouter) -> None:
+    bad = chunk_response(b"bad", rows=1, has_more=False)
+    bad.headers["x-content-sha256"] = hashlib.sha256(b"good").hexdigest()
+    route = respx_mock.get("/audit-logs/export/chunk").mock(return_value=bad)
+
+    with pytest.raises(AuditLogDownloadError, match="checksum mismatch"):
+        client.audit_logs.download(
+            to=BytesIO(),
+            start="2026-06-01T00:00:00Z",
+            end="2026-06-02T00:00:00Z",
+            max_transfer_retries=0,
+        )
+
+    assert route.call_count == 1
+
+
+def test_download_rejects_cursor_cycle(client: Kernel, respx_mock: MockRouter) -> None:
+    respx_mock.get("/audit-logs/export/chunk").mock(
+        side_effect=[
+            chunk_response(b"first", rows=1, has_more=True, next_cursor="a"),
+            chunk_response(b"second", rows=1, has_more=True, next_cursor="b"),
+            chunk_response(b"duplicate", rows=1, has_more=True, next_cursor="a"),
+        ]
+    )
+    destination = BytesIO()
+
+    with pytest.raises(AuditLogDownloadError, match="response repeated X-Next-Cursor header"):
+        client.audit_logs.download(
+            to=destination,
+            start="2026-06-01T00:00:00Z",
+            end="2026-06-02T00:00:00Z",
+        )
+
+    assert destination.getvalue() == b"firstsecond"
 
 
 def test_download_rejects_invalid_cursor_before_writing(client: Kernel, respx_mock: MockRouter) -> None: