fix(progress): don't regress progress when an update is in flight

src/dedalus_mcp/progress.py

@@ -223,6 +223,9 @@ def __init__(
         self._lock = anyio.Lock()
         self._latest: _ProgressState | None = None
         self._last_emitted: _ProgressState | None = None
+        # State taken by ``run`` but not yet delivered (it may be sleeping out the
+        # throttle delay or retrying). It is newer than ``_last_emitted``.
+        self._in_flight: _ProgressState | None = None
         self._last_emit_ns: int = 0
         self._closed = False
         self._pending_updates = 0
@@ -258,8 +261,10 @@ async def set(self, progress: float, *, message: str | None = None, total: float
     async def close(self) -> None:
         async with self._lock:
             if not self._closed:
-                if self._latest is None and self._last_emitted is not None:
+                if self._latest is None and self._in_flight is None and self._last_emitted is not None:
                     # Re-emit the final state to guarantee at-least-once delivery.
+                    # Skip when a newer state is in flight: it will be delivered, and
+                    # re-queueing ``_last_emitted`` would send a stale, lower value after it.
                     self._latest = self._last_emitted
                 self._closed = True
         self._event.set()
@@ -281,6 +286,7 @@ async def run(self) -> None:
                             self._event = anyio.Event()
                             break
                         self._latest = None
+                        self._in_flight = state
                         delay = self._compute_delay(state.timestamp_ns)
                     if delay:
                         await anyio.sleep(delay)
@@ -295,17 +301,18 @@ async def run(self) -> None:
             if not self._drained.is_set():
                 self._drained.set()
 
+    def _newest_state_locked(self) -> _ProgressState | None:
+        """Most recent state: queued, then in flight, then last delivered."""
+        return self._latest or self._in_flight or self._last_emitted
+
     def _current_progress_locked(self) -> float:
-        if self._latest is not None:
-            return self._latest.progress
-        if self._last_emitted is not None:
-            return self._last_emitted.progress
-        return 0.0
+        newest = self._newest_state_locked()
+        return newest.progress if newest is not None else 0.0
 
     def _store_state_locked(
         self, progress: float, *, message: str | None, total_override: float | None
     ) -> tuple[_ProgressState, bool]:
-        previous = self._latest or self._last_emitted
+        previous = self._newest_state_locked()
         if previous and progress < previous.progress:
             raise ValueError(
                 "progress must be monotonically increasing per MCP specification "
@@ -395,6 +402,8 @@ async def _send_with_retry(self, state: _ProgressState) -> None:
             else:
                 async with self._lock:
                     self._last_emitted = state
+                    if self._in_flight is state:
+                        self._in_flight = None
                     self._last_emit_ns = time.monotonic_ns()
                 self._emitted_count += 1
                 self._logger.debug(

tests/test_progress_flow.py

@@ -132,3 +132,190 @@ async def test_progress_throttle_event_emitted() -> None:
 
     assert throttled, "expected throttle telemetry when emit_hz is low"
     assert session.notifications[-1].progress == 2
+
+
+def _assert_monotonic(notifications: list[RecordedNotification]) -> None:
+    values = [n.progress for n in notifications]
+    assert values == sorted(values), f"progress went backwards: {values}"
+
+
+@pytest.mark.anyio
+async def test_close_during_throttled_send_does_not_regress() -> None:
+    """Closing while a throttled update is in flight must not re-send an older value."""
+    session = FakeSession()
+    closes: list[ProgressCloseEvent] = []
+    telemetry = ProgressTelemetry(on_close=lambda evt: closes.append(evt))
+
+    async with with_request_context(token="tok", session=session):
+        async with progress(total=4, telemetry=telemetry, config=ProgressConfig(emit_hz=20)) as tracker:
+            await tracker.advance(1, "one")
+            await anyio.sleep(0.005)  # first update delivered immediately
+            await tracker.set(4, message="done")  # throttled
+            await anyio.sleep(0.001)  # emitter picks it up; still in flight at close
+
+    _assert_monotonic(session.notifications)
+    assert session.notifications[-1].progress == 4
+    assert closes
+    assert closes[-1].final_progress == 4
+    assert closes[-1].final_message == "done"
+
+
+@pytest.mark.anyio
+async def test_advance_uses_in_flight_state_as_base() -> None:
+    """``advance`` must build on an update that is in flight, not the last delivered one."""
+    session = FakeSession()
+
+    async with with_request_context(token="tok", session=session):
+        async with progress(total=10, config=ProgressConfig(emit_hz=20)) as tracker:
+            await tracker.advance(1)
+            await anyio.sleep(0.005)  # first update delivered immediately
+            await tracker.set(3)  # throttled
+            await anyio.sleep(0.001)  # emitter picks it up and waits out the throttle
+            assert await tracker.advance(1) == 4
+
+    _assert_monotonic(session.notifications)
+    assert session.notifications[-1].progress == 4
+
+
+class HeldSession(FakeSession):
+    """FakeSession that holds each send until the test releases it.
+
+    While a send is held, the emitter keeps that state in flight, so tests can act
+    inside the race window without relying on sleeps. ``fail_attempts`` lists send
+    attempts (1-based, across all sends) that raise after being released.
+    """
+
+    def __init__(self, *, fail_attempts: tuple[int, ...] = ()) -> None:
+        super().__init__()
+        self.attempts = 0
+        self._fail_attempts = fail_attempts
+        self._holding = True
+        self._started = anyio.Event()
+        self._gate = anyio.Event()
+
+    async def send_progress_notification(
+        self,
+        progress_token: str | int,
+        progress_value: float,
+        *,
+        total: float | None = None,
+        message: str | None = None,
+    ) -> None:
+        self.attempts += 1
+        attempt = self.attempts
+        if self._holding:
+            gate = self._gate
+            self._started.set()
+            await gate.wait()
+        if attempt in self._fail_attempts:
+            raise RuntimeError("simulated send failure")
+        await super().send_progress_notification(progress_token, progress_value, total=total, message=message)
+
+    async def wait_for_send(self) -> None:
+        """Wait until the emitter is blocked inside a send."""
+        await self._started.wait()
+
+    def release(self) -> None:
+        """Let the held send finish; the next send will be held again."""
+        gate = self._gate
+        self._started = anyio.Event()
+        self._gate = anyio.Event()
+        gate.set()
+
+    def stop_holding(self) -> None:
+        """Release the held send, if any, and let all later sends through."""
+        self._holding = False
+        self.release()
+
+
+_HELD_CONFIG = ProgressConfig(emit_hz=0, retry_backoff=(0.0, 0.0))
+
+
+async def _stop_holding_once_blocked(session: HeldSession) -> None:
+    """Release sends only once ``close()`` is waiting, so it runs while a state is in flight."""
+    await anyio.wait_all_tasks_blocked()
+    session.stop_holding()
+
+
+@pytest.mark.anyio
+async def test_close_while_send_in_flight_does_not_regress() -> None:
+    """Deterministic version of the close race: the send of 4 is held while close() runs."""
+    session = HeldSession()
+    closes: list[ProgressCloseEvent] = []
+    telemetry = ProgressTelemetry(on_close=lambda evt: closes.append(evt))
+
+    async with with_request_context(token="tok", session=session), anyio.create_task_group() as tg:
+        async with progress(total=4, telemetry=telemetry, config=_HELD_CONFIG) as tracker:
+            await tracker.advance(1, "one")
+            await session.wait_for_send()
+            session.release()  # 1 delivered
+            await tracker.set(4, message="done")
+            await session.wait_for_send()  # 4 is in flight
+            tg.start_soon(_stop_holding_once_blocked, session)
+
+    assert [n.progress for n in session.notifications] == [1, 4]
+    assert closes[-1].final_progress == 4
+    assert closes[-1].final_message == "done"
+
+
+@pytest.mark.anyio
+async def test_advance_builds_on_held_in_flight_send() -> None:
+    """Deterministic version of the advance race: the send of 3 is held while advance() runs."""
+    session = HeldSession()
+
+    async with with_request_context(token="tok", session=session):
+        async with progress(total=10, config=_HELD_CONFIG) as tracker:
+            await tracker.advance(1)
+            await session.wait_for_send()
+            session.release()  # 1 delivered
+            await tracker.set(3)
+            await session.wait_for_send()  # 3 is in flight
+            try:
+                assert await tracker.advance(1) == 4
+            finally:
+                session.stop_holding()
+
+    _assert_monotonic(session.notifications)
+    assert session.notifications[-1].progress == 4
+
+
+@pytest.mark.anyio
+async def test_close_during_retry_does_not_regress() -> None:
+    """A state being retried after a failed send is still in flight; close() must not re-send an older one."""
+    session = HeldSession(fail_attempts=(2,))
+
+    async with with_request_context(token="tok", session=session), anyio.create_task_group() as tg:
+        async with progress(total=4, config=_HELD_CONFIG) as tracker:
+            await tracker.set(1)
+            await session.wait_for_send()
+            session.release()  # attempt 1: 1 delivered
+            await tracker.set(4)
+            await session.wait_for_send()
+            session.release()  # attempt 2: sending 4 fails
+            await session.wait_for_send()  # attempt 3: 4 is being retried
+            tg.start_soon(_stop_holding_once_blocked, session)
+
+    assert [n.progress for n in session.notifications] == [1, 4]
+    assert session.attempts == 3
+
+
+@pytest.mark.anyio
+async def test_set_below_in_flight_value_raises() -> None:
+    """Monotonicity is checked against the in-flight state, not just the last delivered one."""
+    session = HeldSession()
+
+    async with with_request_context(token="tok", session=session):
+        async with progress(total=10, config=_HELD_CONFIG) as tracker:
+            await tracker.set(1)
+            await session.wait_for_send()
+            session.release()  # 1 delivered
+            await tracker.set(3)
+            await session.wait_for_send()  # 3 is in flight
+            try:
+                with pytest.raises(ValueError):
+                    await tracker.set(2)
+            finally:
+                session.stop_holding()
+
+    _assert_monotonic(session.notifications)
+    assert session.notifications[-1].progress == 3