fix: preserve uncertainty across append retries

src/s2_sdk/__init__.py

@@ -5,6 +5,7 @@
 from s2_sdk._batching import append_inputs, append_record_batches
 from s2_sdk._exceptions import (
     AppendConditionError,
+    AppendIndefiniteFailureError,
     FencingTokenMismatchError,
     ReadUnwrittenError,
     S2ClientError,
@@ -133,6 +134,7 @@
     "S2ClientError",
     "S2ServerError",
     "AppendConditionError",
+    "AppendIndefiniteFailureError",
     "FencingTokenMismatchError",
     "SeqNumMismatchError",
     "ReadUnwrittenError",

src/s2_sdk/_exceptions.py

@@ -69,6 +69,23 @@ class ServerDrainingError(S2ServerError):
     """The server is draining, so the associated connection should be retired."""
 
 
+class AppendIndefiniteFailureError(S2Error):
+    """The final append attempt failed definitively, but an earlier attempt may have
+    taken effect, so the entire append operation is indefinite.
+
+    Attributes:
+        final_attempt_error: The definite error raised by the final attempt.
+    """
+
+    def __init__(self, final_attempt_error: Exception):
+        self.final_attempt_error = final_attempt_error
+        super().__init__(
+            "append may have taken effect in an earlier attempt; "
+            f"final attempt failed: {final_attempt_error}"
+        )
+        self.__cause__ = final_attempt_error
+
+
 class AppendConditionError(S2ServerError):
     """Append condition was not met."""
 

src/s2_sdk/_ops.py

@@ -976,6 +976,7 @@ def __init__(
             max_retries=retry._max_retries(),
             min_base_delay=retry.min_base_delay.total_seconds(),
             max_base_delay=retry.max_base_delay.total_seconds(),
+            track_append_uncertainty=True,
         )
 
     def __repr__(self) -> str:

src/s2_sdk/_retrier.py

@@ -7,6 +7,7 @@
 from typing import Callable
 
 from s2_sdk._exceptions import (
+    AppendIndefiniteFailureError,
     ConnectError,
     ConnectionRetiredError,
     S2ServerError,
@@ -29,20 +30,25 @@ def __init__(
         max_retries: int,
         min_base_delay: float = 0.1,
         max_base_delay: float = 1.0,
+        track_append_uncertainty: bool = False,
     ):
         self.should_retry_on = should_retry_on
         self.max_retries = max_retries
         self.min_base_delay = min_base_delay
         self.max_base_delay = max_base_delay
+        self.track_append_uncertainty = track_append_uncertainty
 
     async def __call__(self, f: Callable, *args, **kwargs):
         max_retries = self.max_retries
         attempt = 0
+        prior_uncertainty = False
         while True:
             try:
                 return await f(*args, **kwargs)
             except Exception as e:
                 if attempt < max_retries and self.should_retry_on(e):
+                    if self.track_append_uncertainty and not has_no_side_effects(e):
+                        prior_uncertainty = True
                     delay = compute_backoff(
                         attempt,
                         min_base_delay=self.min_base_delay,
@@ -66,6 +72,10 @@ async def __call__(self, f: Callable, *args, **kwargs):
                         self.should_retry_on(e),
                         attempt >= max_retries,
                     )
+                    if self.track_append_uncertainty:
+                        wrapped = with_prior_uncertainty(e, prior_uncertainty)
+                        if wrapped is not e:
+                            raise wrapped from e
                     raise e
 
 
@@ -160,6 +170,8 @@ def http_retry_on(e: Exception) -> bool:
 
 
 def has_no_side_effects(e: Exception) -> bool:
+    if isinstance(e, AppendIndefiniteFailureError):
+        return False
     if requires_reconnect(e):
         return True
     if isinstance(e, S2ServerError):
@@ -174,3 +186,25 @@ def has_no_side_effects(e: Exception) -> bool:
             cause = cause.__cause__
         return False
     return False
+
+
+def with_prior_uncertainty(e: Exception, prior_uncertainty: bool) -> Exception:
+    """Wrap a definite final error if an earlier attempt may have taken effect.
+
+    Already indefinite errors are returned unchanged.
+    """
+    if prior_uncertainty and has_no_side_effects(e):
+        return AppendIndefiniteFailureError(e)
+    return e
+
+
+def attempt_may_have_side_effects(
+    e: Exception, frame_signal: FrameSignal | None
+) -> bool:
+    """Whether this attempt may have taken effect, accounting for unsent request data.
+
+    Without a frame signal, assume the request may have been sent.
+    """
+    return not has_no_side_effects(e) and (
+        frame_signal is None or frame_signal.is_signalled()
+    )

src/s2_sdk/_s2s/_append_session.py

@@ -19,9 +19,11 @@
 from s2_sdk._retrier import (
     AdvisedReconnectLimiter,
     Attempt,
+    attempt_may_have_side_effects,
     compute_backoff,
     is_safe_to_retry_session,
     requires_reconnect,
+    with_prior_uncertainty,
 )
 from s2_sdk._s2s import _stream_records_path
 from s2_sdk._s2s._protocol import (
@@ -52,6 +54,9 @@ class _InflightInput:
     num_records: int
     encoded: bytes
     ack_deadline: float | None = None
+    # Set once an attempt carrying this input failed in a way that may have
+    # taken effect and was retried.
+    prior_uncertainty: bool = False
 
 
 @dataclass(slots=True)
@@ -146,6 +151,9 @@ async def retrying_inner():
                     bool(session_state.inflight_inputs),
                     frame_signal,
                 ):
+                    if attempt_may_have_side_effects(e, frame_signal):
+                        for inflight in session_state.inflight_inputs:
+                            inflight.prior_uncertainty = True
                     backoff = compute_backoff(
                         attempt.value,
                         min_base_delay=min_base_delay,
@@ -165,7 +173,16 @@ async def retrying_inner():
                         e,
                         attempt.value >= max_retries,
                     )
-                    raise
+                    wrapped = with_prior_uncertainty(
+                        e,
+                        any(
+                            inflight.prior_uncertainty
+                            for inflight in session_state.inflight_inputs
+                        ),
+                    )
+                    if wrapped is e:
+                        raise
+                    raise wrapped from e
         await ack_queue.put(None)
 
     async with asyncio.TaskGroup() as tg:

tests/test_append_session.py

@@ -4,9 +4,12 @@
 from typing import Any, cast
 from unittest.mock import patch
 
+import pytest
+
 import s2_sdk._generated.s2.v1.s2_pb2 as pb
 import s2_sdk._s2s._append_session as append_session
 from s2_sdk._client import HttpClient
+from s2_sdk._exceptions import AppendIndefiniteFailureError, S2ServerError
 from s2_sdk._s2s._protocol import Message
 from s2_sdk._types import AppendInput, Compression, Record, Retry
 
@@ -93,3 +96,63 @@ async def inputs() -> AsyncIterator[AppendInput]:
     assert [ack.end.seq_num for ack in acks] == [1, 2]
     assert len(responses) == 2
     assert responses[0].retired
+
+
+@pytest.mark.parametrize(
+    ("first_code", "first_status", "want_wrapped"),
+    [
+        ("unavailable", 503, True),
+        ("rate_limited", 429, False),
+    ],
+)
+async def test_terminal_definite_error_preserves_prior_uncertainty(
+    first_code: str, first_status: int, want_wrapped: bool
+) -> None:
+    errors = [
+        S2ServerError(code=first_code, message=first_code, status_code=first_status),
+        S2ServerError(code="rate_limited", message="rate_limited", status_code=429),
+    ]
+    final = errors[-1]
+
+    class _Client:
+        @asynccontextmanager
+        async def streaming_request(
+            self,
+            *args: Any,
+            content: AsyncGenerator[bytes, None] | None = None,
+            **kwargs: Any,
+        ) -> AsyncGenerator[_Response, None]:
+            assert content is not None
+            try:
+                # Consume one input so it becomes inflight, then fail.
+                await content.__anext__()
+                raise errors.pop(0)
+            finally:
+                await content.aclose()
+            yield _Response(())  # pragma: no cover
+
+    async def inputs() -> AsyncIterator[AppendInput]:
+        yield AppendInput(records=[Record(body=b"a")])
+
+    with (
+        patch.object(append_session, "compute_backoff", new=lambda *a, **k: 0.0),
+        pytest.raises(BaseException) as exc_info,
+    ):
+        async for _ in append_session.run_append_session(
+            cast(HttpClient, _Client()),
+            "stream",
+            inputs(),
+            Retry(max_attempts=2),
+            Compression.NONE,
+            ack_timeout=1.0,
+        ):
+            pass
+
+    err: BaseException = exc_info.value
+    while isinstance(err, BaseExceptionGroup):
+        err = err.exceptions[0]
+    if want_wrapped:
+        assert isinstance(err, AppendIndefiniteFailureError)
+        assert err.final_attempt_error is final
+    else:
+        assert err is final

tests/test_retrier.py

@@ -2,7 +2,81 @@
 
 import pytest
 
-from s2_sdk._retrier import compute_backoff
+from s2_sdk._exceptions import AppendIndefiniteFailureError, S2ServerError
+from s2_sdk._retrier import (
+    Retrier,
+    compute_backoff,
+    has_no_side_effects,
+    is_safe_to_retry_unary,
+    with_prior_uncertainty,
+)
+from s2_sdk._types import AppendRetryPolicy
+
+
+def _server_error(status_code: int, code: str) -> S2ServerError:
+    return S2ServerError(code=code, message=code, status_code=status_code)
+
+
+class TestWithPriorUncertainty:
+    def test_definite_error_unchanged_without_prior_uncertainty(self):
+        e = _server_error(429, "rate_limited")
+        assert with_prior_uncertainty(e, False) is e
+
+    def test_indefinite_error_not_wrapped(self):
+        e = _server_error(503, "unavailable")
+        assert with_prior_uncertainty(e, True) is e
+
+    def test_definite_error_wrapped(self):
+        e = _server_error(429, "rate_limited")
+        wrapped = with_prior_uncertainty(e, True)
+        assert isinstance(wrapped, AppendIndefiniteFailureError)
+        assert wrapped.final_attempt_error is e
+        assert wrapped.__cause__ is e
+        assert not has_no_side_effects(wrapped)
+        assert with_prior_uncertainty(wrapped, True) is wrapped
+
+
+class TestRetrierAppendUncertainty:
+    def _retrier(self) -> Retrier:
+        return Retrier(
+            should_retry_on=lambda e: is_safe_to_retry_unary(e, AppendRetryPolicy.ALL),
+            max_retries=2,
+            min_base_delay=0.001,
+            max_base_delay=0.001,
+            track_append_uncertainty=True,
+        )
+
+    async def _run(self, responses: list[Exception | str]):
+        it = iter(responses)
+
+        async def f():
+            r = next(it)
+            if isinstance(r, Exception):
+                raise r
+            return r
+
+        return await self._retrier()(f)
+
+    async def test_indefinite_then_definite_is_wrapped(self):
+        final = _server_error(429, "rate_limited")
+        with pytest.raises(AppendIndefiniteFailureError) as exc_info:
+            await self._run([_server_error(503, "unavailable"), final, final])
+        assert exc_info.value.final_attempt_error is final
+
+    async def test_definite_then_definite_is_not_wrapped(self):
+        final = _server_error(429, "rate_limited")
+        with pytest.raises(S2ServerError) as exc_info:
+            await self._run([_server_error(429, "rate_limited"), final, final])
+        assert exc_info.value is final
+
+    async def test_indefinite_then_indefinite_is_not_wrapped(self):
+        final = _server_error(503, "unavailable")
+        with pytest.raises(S2ServerError) as exc_info:
+            await self._run([_server_error(503, "unavailable"), final, final])
+        assert exc_info.value is final
+
+    async def test_success_after_indefinite(self):
+        assert await self._run([_server_error(503, "unavailable"), "ok"]) == "ok"
 
 
 class TestComputeBackoff: