diff --git a/src/dedalus_mcp/progress.py b/src/dedalus_mcp/progress.py index 28158bd..d64fdea 100644 --- a/src/dedalus_mcp/progress.py +++ b/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( diff --git a/tests/test_progress_flow.py b/tests/test_progress_flow.py index 28061dd..b247e8d 100644 --- a/tests/test_progress_flow.py +++ b/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