fix(retry): correct attempt-index off-by-one between wait strategies (#734)

## Summary

Fixes an off-by-one error in retry wait-strategy indexing.

## Problem

`_ComposableRetryPolicy.next()` forwarded the runtime's 1-indexed
`attempts` value (which always starts at 1 on the first failure, see
`runtime/control_loop/reduce.py: failures = this_execution.attempts +
1`) straight to the configured wait strategy.

Every attempt-indexed wait strategy (`wait_chain`, `wait_exponential`,
`wait_incrementing`, `wait_exponential_jitter`,
`wait_random_exponential`) is documented and unit-tested as 0-indexed,
where index 0 is the delay before the first retry.

The mismatch meant every wait strategy silently skipped its first
configured delay and used the second delay for the first retry, third
for the second, and so on. For `wait_chain` specifically, this meant the
first strategy in the chain was never used at all.

## Solution

Subtract 1 from `attempts` before passing it to the wait strategy, so it
receives the 0-indexed value its contract expects.

Stop conditions are unaffected and continue to receive the original
1-indexed `attempts`, which is the correct convention for
`stop_after_attempt`.

The existing unit tests for `.next()` were also updated because they
called it directly with attempts starting at 0, an assumption that did
not match how the runtime actually invokes it. A regression test was
added that reproduces the real retry-loop calling convention end-to-end
and fails with the previous behavior.

## Tests

- Full test suite: 857 passed
- Ruff formatting/lint checks passed
- Type checking passed

Fixes #733

---------

Co-authored-by: Adrian Lyjak <adrianlyjak@gmail.com>

.changeset/retry-wait-indexing.md

@@ -0,0 +1,5 @@
+---
+"llama-index-workflows": patch
+---
+
+Fix an off-by-one that made retry wait strategies skip their first configured delay, so every retry delay is now one step shorter than before.

docs/src/content/docs/llamaagents/workflows/retry_steps.md

@@ -162,6 +162,7 @@ def next(
 ```
 
 Return the number of seconds to wait before retrying, or `None` to stop.
+`attempts` counts failures so far, starting at 1 for the first failure.
 The optional `seed` is used by durable runtimes to make jitter deterministic during replay.
 
 For example, this policy only retries on Fridays:

packages/llama-index-workflows/src/workflows/retry_policy.py

@@ -90,7 +90,8 @@ class RetryPolicy(Protocol):
 
         Args:
             elapsed_time: Seconds since the first failure.
-            attempts: Number of attempts made so far.
+            attempts: Number of failures so far, starting at 1 for the
+                first failure.
             error: The last exception encountered.
             seed: Optional RNG seed for deterministic jitter (DBOS replay).
 
@@ -106,7 +107,11 @@ class RetryCondition(Protocol):
 
 
 class WaitStrategy(Protocol):
-    """Compute the delay in seconds before the next retry attempt."""
+    """Compute the delay in seconds before the next retry attempt.
+
+    ``attempts`` starts at 0 for the first retry. ``RetryPolicy.next`` counts
+    failures from 1 and subtracts one before calling the wait strategy.
+    """
 
     def __call__(self, attempts: int, *, seed: int | None = None) -> float: ...
 
@@ -463,7 +468,8 @@ class wait_exponential(_WaitStrategyBase):
     """
     Wait with exponentially increasing delays, clamped between ``min`` and ``max``.
 
-    The delay for attempt ``n`` is ``multiplier * exp_base**n`` before clamping.
+    The delay for retry ``n`` is ``multiplier * exp_base**n`` before clamping,
+    with ``n=0`` for the first retry.
 
     Examples:
         ```python
@@ -494,8 +500,8 @@ class wait_incrementing(_WaitStrategyBase):
     """
     Wait an incrementally larger amount after each attempt.
 
-    The delay starts at ``start`` and increases by ``increment`` on each retry,
-    capped by ``max`` and never going below zero.
+    The delay is ``start`` for the first retry and increases by ``increment``
+    on each subsequent retry, capped by ``max`` and never going below zero.
 
     Examples:
         ```python
@@ -630,10 +636,11 @@ def wait_full_jitter(
 
 class wait_chain(_WaitStrategyBase):
     """
-    Use a different wait strategy for each attempt in order.
+    Use a different wait strategy for each retry in order, starting with the
+    first.
 
     After the provided strategies are exhausted, the last strategy is reused
-    for all subsequent attempts.
+    for all subsequent retries.
 
     Examples:
         ```python
@@ -832,7 +839,12 @@ class _ComposableRetryPolicy:
         if self.retry is not None and not self.retry(error):
             return None
 
-        delay = self.wait(attempts, seed=seed)
+        # `attempts` counts failures from 1 (see runtime/control_loop/
+        # reduce.py: `failures = this_execution.attempts + 1`) while wait
+        # strategies count retries from 0, so subtract one. The clamp keeps
+        # a direct caller passing 0 on the first delay instead of index -1,
+        # which wait_chain would map to its last strategy.
+        delay = self.wait(max(attempts - 1, 0), seed=seed)
         if self.stop(attempts, elapsed_time, upcoming_sleep=delay):
             return None
         return delay

packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py

@@ -33,6 +33,7 @@ from workflows.retry_policy import (
     ExponentialBackoffRetryPolicy,
     retry_policy,
     stop_after_attempt,
+    wait_chain,
     wait_exponential_jitter,
     wait_fixed,
 )
@@ -733,6 +734,42 @@ def test_step_worker_failed_retry_preserves_delay(base_state: BrokerState) -> No
     assert len(new_state.workers["test_step"].in_progress) == 0
 
 
+@pytest.mark.parametrize(
+    "prior_attempts, expected_delay",
+    [(0, 1.0), (1, 2.0), (2, 5.0)],
+)
+def test_step_worker_failed_uses_attempt_indexed_delay(
+    base_state: BrokerState, prior_attempts: int, expected_delay: float
+) -> None:
+    """The Nth failure re-queues with the Nth configured delay.
+
+    A constant delay can't see an indexing off-by-one, so this uses
+    wait_chain with three distinct values and checks the absolute
+    not_before for each failure.
+    """
+    base_state.workers["test_step"].config.retry_policy = retry_policy(
+        wait=wait_chain(wait_fixed(1.0), wait_fixed(2.0), wait_fixed(5.0)),
+        stop=stop_after_attempt(4),
+    )
+    event = MyTestEvent(value=42)
+    add_worker(base_state, event)
+    base_state.workers["test_step"].in_progress[0].attempts = prior_attempts
+
+    tick: TickStepResult = TickStepResult(
+        step_id=StepId.root("test_step"),
+        worker_id=0,
+        event=event,
+        result=[StepWorkerFailed(exception=ValueError("test"), failed_at=110.0)],
+    )
+
+    new_state, _ = _process_step_result_tick(tick, base_state, now_seconds=110.0)
+
+    queue = new_state.workers["test_step"].queue
+    assert len(queue) == 1
+    assert queue[0].not_before == 110.0 + expected_delay
+    assert queue[0].attempts == prior_attempts + 1
+
+
 def test_step_worker_failed_retry_preserves_first_attempt_at(
     base_state: BrokerState,
 ) -> None:

packages/llama-index-workflows/tests/test_retry_policy.py

@@ -438,9 +438,11 @@ def test_retry_policy_with_wait_strategy() -> None:
         wait=wait_exponential(multiplier=1.0, exp_base=2.0, max=100.0, min=0.0),
         stop=stop_after_attempt(max_attempt_number=5),
     )
-    assert p.next(0.0, 0, Exception()) == 1.0
-    assert p.next(0.0, 1, Exception()) == 2.0
-    assert p.next(0.0, 2, Exception()) == 4.0
+    # `attempts` mirrors the real call site: 1 on the first failure, 2 on the
+    # second, etc. (see runtime/control_loop/reduce.py: failures = attempts + 1).
+    assert p.next(0.0, 1, Exception()) == 1.0
+    assert p.next(0.0, 2, Exception()) == 2.0
+    assert p.next(0.0, 3, Exception()) == 4.0
 
 
 def test_retry_policy_with_random_exponential_wait() -> None:
@@ -469,8 +471,9 @@ def test_retry_policy_stop_before_delay_stops_using_next_sleep() -> None:
         wait=wait_incrementing(start=1.0, increment=1.0, max=10.0),
         stop=stop_before_delay(5.0),
     )
-    assert p.next(2.0, 1, Exception("retry")) == 2.0
-    assert p.next(2.0, 2, Exception("retry")) is None
+    assert p.next(2.0, 1, Exception("retry")) == 1.0
+    assert p.next(2.0, 2, Exception("retry")) == 2.0
+    assert p.next(2.0, 3, Exception("retry")) is None
 
 
 def test_retry_policy_seed_forwarded() -> None:
@@ -496,11 +499,41 @@ def test_retry_policy_all_three_composed() -> None:
         stop=stop_after_attempt(max_attempt_number=3),
     )
     err = ConnectionError("refused")
-    assert p.next(0.0, 0, err) == 0.5
-    assert p.next(0.0, 1, err) == 1.5
-    assert p.next(0.0, 2, err) == 4.5
+    # `attempts` counts failures so far, matching how the runtime actually
+    # calls `.next()` (1 on the first failure, never 0 — see
+    # runtime/control_loop/reduce.py: `failures = this_execution.attempts + 1`).
+    assert p.next(0.0, 1, err) == 0.5
+    assert p.next(0.0, 2, err) == 1.5
     assert p.next(0.0, 3, err) is None
-    assert p.next(0.0, 0, ValueError("bad")) is None
+    assert p.next(0.0, 1, ValueError("bad")) is None
+
+
+def test_retry_policy_wait_indexing_matches_real_call_site() -> None:
+    """Regression test for an off-by-one between wait strategies and the
+    runtime's actual calling convention.
+
+    The runtime always calls ``RetryPolicy.next()`` with ``attempts`` counting
+    completed failures starting at 1 (see
+    ``runtime/control_loop/reduce.py``: ``failures = this_execution.attempts +
+    1``); it is never 0. Every attempt-indexed wait strategy (``wait_chain``,
+    ``wait_exponential``, ``wait_incrementing``, ...) is documented and
+    unit-tested as 0-indexed: the first configured delay corresponds to index
+    0. Before the fix, ``_ComposableRetryPolicy.next()`` forwarded the
+    1-indexed ``attempts`` straight to the wait strategy, silently skipping
+    each strategy's first configured delay and shifting every later delay one
+    step ahead of schedule.
+    """
+    p = retry_policy(
+        wait=wait_chain(wait_fixed(1.0), wait_fixed(2.0), wait_fixed(5.0)),
+        stop=stop_after_attempt(4),
+    )
+    err = Exception("boom")
+
+    # Simulate the real retry loop: `attempts` starts at 1 on the first
+    # failure and increments by 1 on each subsequent failure.
+    delays = [p.next(0.0, attempts, err) for attempts in range(1, 4)]
+
+    assert delays == [1.0, 2.0, 5.0]
 
 
 def test_retry_policy_with_operator_composition() -> None:
@@ -561,10 +594,10 @@ def test_ExponentialBackoffRetryPolicy_next_basic() -> None:
     )
     assert type(p).__name__ == "_ComposableRetryPolicy"
     err = Exception()
-    assert p.next(elapsed_time=0.0, attempts=0, error=err) == 1.0
-    assert p.next(elapsed_time=0.0, attempts=1, error=err) == 2.0
-    assert p.next(elapsed_time=0.0, attempts=2, error=err) == 4.0
-    assert p.next(elapsed_time=0.0, attempts=3, error=err) == 8.0
+    assert p.next(elapsed_time=0.0, attempts=1, error=err) == 1.0
+    assert p.next(elapsed_time=0.0, attempts=2, error=err) == 2.0
+    assert p.next(elapsed_time=0.0, attempts=3, error=err) == 4.0
+    assert p.next(elapsed_time=0.0, attempts=4, error=err) == 8.0
 
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
@@ -572,7 +605,7 @@ def test_ExponentialBackoffRetryPolicy_max_delay_cap() -> None:
     p = ExponentialBackoffRetryPolicy(
         initial_delay=1.0, multiplier=10.0, max_delay=50.0, jitter=False
     )
-    assert p.next(elapsed_time=0.0, attempts=2, error=Exception()) == 50.0
+    assert p.next(elapsed_time=0.0, attempts=3, error=Exception()) == 50.0
     assert p.next(elapsed_time=0.0, attempts=5, error=Exception()) is None
 
 
@@ -624,10 +657,10 @@ def test_ExponentialBackoffRetryPolicy_no_jitter() -> None:
         initial_delay=0.5, multiplier=3.0, max_delay=100.0, jitter=False
     )
     err = Exception()
-    assert p.next(elapsed_time=0.0, attempts=0, error=err) == 0.5
-    assert p.next(elapsed_time=0.0, attempts=1, error=err) == 1.5
-    assert p.next(elapsed_time=0.0, attempts=2, error=err) == 4.5
-    assert p.next(elapsed_time=0.0, attempts=3, error=err) == 13.5
+    assert p.next(elapsed_time=0.0, attempts=1, error=err) == 0.5
+    assert p.next(elapsed_time=0.0, attempts=2, error=err) == 1.5
+    assert p.next(elapsed_time=0.0, attempts=3, error=err) == 4.5
+    assert p.next(elapsed_time=0.0, attempts=4, error=err) == 13.5
 
 
 # ---------------------------------------------------------------------------

packages/llama-index-workflows/tests/test_retry_tenacity_conformance.py

@@ -7,9 +7,10 @@ import inspect
 import types as builtin_types
 import typing
 from collections.abc import Callable
-from typing import Protocol, Union, cast
+from typing import Any, Protocol, Union, cast
 
 import pytest
+import tenacity
 from tenacity import retry_all as tenacity_retry_all
 from tenacity import retry_any as tenacity_retry_any
 from tenacity import retry_if_exception as tenacity_retry_if_exception
@@ -56,6 +57,7 @@ from workflows.retry_policy import (
     retry_if_not_exception_message,
     retry_if_not_exception_type,
     retry_never,
+    retry_policy,
     retry_unless_exception_type,
     stop_after_attempt,
     stop_after_delay,
@@ -205,3 +207,88 @@ def test_wait_full_jitter_alias_matches_wait_random_exponential_signature() -> N
     assert _parameter_types(wait_full_jitter) == _parameter_types(
         wait_random_exponential
     )
+
+
+# ---------------------------------------------------------------------------
+# Value conformance: same config in, same delay sequence out.
+#
+# The signature tests above catch API drift but not behavioral drift. Each
+# case drives our policy the way the runtime does (attempts counts failures
+# from 1) and tenacity's wait strategy the way its retry loop does
+# (attempt_number from 1), and requires identical delays. Deterministic
+# strategies only, since tenacity's jitter uses global random state.
+# ---------------------------------------------------------------------------
+
+
+def _tenacity_delays(strategy: Callable[..., float], n: int) -> list[float]:
+    """Delay sequence tenacity's retry loop would produce for n failures."""
+    delays = []
+    for attempt_number in range(1, n + 1):
+        state = tenacity.RetryCallState(cast(Any, None), None, (), {})
+        state.attempt_number = attempt_number
+        delays.append(strategy(state))
+    return delays
+
+
+def _our_delays(wait: object, n: int) -> list[float | None]:
+    """Delay sequence our runtime would produce for n failures."""
+    policy = retry_policy(wait=cast(Any, wait), stop=stop_never())
+    err = Exception("boom")
+    return [policy.next(0.0, attempts, err) for attempts in range(1, n + 1)]
+
+
+VALUE_CONFORMANCE_CASES: list[tuple[str, object, Callable[..., float]]] = [
+    ("wait_fixed", wait_fixed(1.5), tenacity_wait_fixed(1.5)),
+    ("wait_none", wait_none(), tenacity_wait_none()),
+    (
+        "wait_exponential",
+        wait_exponential(multiplier=1.0, exp_base=2.0, max=60.0, min=0.0),
+        tenacity_wait_exponential(multiplier=1.0, exp_base=2.0, max=60.0, min=0.0),
+    ),
+    (
+        "wait_exponential_capped",
+        wait_exponential(multiplier=0.5, exp_base=3.0, max=10.0, min=1.0),
+        tenacity_wait_exponential(multiplier=0.5, exp_base=3.0, max=10.0, min=1.0),
+    ),
+    (
+        "wait_incrementing",
+        wait_incrementing(start=1.0, increment=2.0, max=10.0),
+        tenacity_wait_incrementing(start=1.0, increment=2.0, max=10.0),
+    ),
+    (
+        "wait_exponential_jitter_zero_jitter",
+        wait_exponential_jitter(initial=1.0, exp_base=2.0, max=100.0, jitter=0.0),
+        tenacity_wait_exponential_jitter(
+            initial=1.0, exp_base=2.0, max=100.0, jitter=0.0
+        ),
+    ),
+    (
+        "wait_chain",
+        wait_chain(wait_fixed(1.0), wait_fixed(2.0), wait_fixed(5.0)),
+        tenacity_wait_chain(
+            tenacity_wait_fixed(1.0),
+            tenacity_wait_fixed(2.0),
+            tenacity_wait_fixed(5.0),
+        ),
+    ),
+    (
+        "wait_combine",
+        wait_combine(wait_fixed(1.0), wait_incrementing(start=0.0, increment=1.0)),
+        tenacity_wait_combine(
+            tenacity_wait_fixed(1.0),
+            tenacity_wait_incrementing(start=0.0, increment=1.0),
+        ),
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "wait, tenacity_wait",
+    [(c[1], c[2]) for c in VALUE_CONFORMANCE_CASES],
+    ids=[c[0] for c in VALUE_CONFORMANCE_CASES],
+)
+def test_retry_policy_delays_match_tenacity(
+    wait: object, tenacity_wait: Callable[..., float]
+) -> None:
+    n = 6
+    assert _our_delays(wait, n) == _tenacity_delays(tenacity_wait, n)