diff --git a/EvoScientist/cli/async_notifier.py b/EvoScientist/cli/async_notifier.py index 16cec23..19d705f 100644 --- a/EvoScientist/cli/async_notifier.py +++ b/EvoScientist/cli/async_notifier.py @@ -24,6 +24,12 @@ TERMINAL_STATUSES: Final = frozenset({"success", "error", "timeout", "interrupte Cancel operations transition runs into ``interrupted`` (not ``cancelled``). """ +# How many times the watcher will re-join the SSE stream when it closes +# cleanly but ``runs.get`` reports the run is still alive (typical cause: +# HTTP keep-alive timeout on long static periods). Bounded to prevent an +# unbounded loop if the server permanently misreports status. +_MAX_RECONNECT_ATTEMPTS: Final = 10 + @dataclass(frozen=True) class AsyncTaskNotification: @@ -132,62 +138,126 @@ async def watch_run_and_notify( ) -> None: """Subscribe to a run's event stream; enqueue notification when it terminates. - Status detection strategy: - 1. Watch for an explicit ``event="error"`` SSE part — langgraph dev - emits one when the run fails. This is authoritative, in-band, and - has no timing race against the server-side run-state writeback. - 2. On clean stream exit with no error event → ``"success"``. - 3. On stream exception → fall back to ``client.runs.get`` (best-effort). - Non-terminal fallback statuses (``pending`` / ``running``) are - dropped — the run is still alive, no notification is enqueued. + Status detection strategy (priority order): - Reading the in-band ``event="error"`` SSE part instead of polling - ``runs.get`` after every clean close avoids a race where the server-side - terminal status hasn't been written by the time the stream closes. + 1. **In-band ``event="error"`` SSE part** — authoritative error signal + from langgraph dev, no race against server-side state writeback. + 2. **Server-side state via ``runs.get``** — invoked after the stream + closes (cleanly or with exception) to verify the run is actually + done. Required because SSE long-poll can close on HTTP keep-alive + timeout while the run is still running, which would otherwise be + misread as ``"success"`` (observed in production with long-running + literature search tasks under concurrency). + 3. **Re-join loop** — if ``runs.get`` reports ``pending`` / ``running``, + the run is alive but we lost the stream; re-join up to + ``_MAX_RECONNECT_ATTEMPTS`` times before giving up. + + The previous implementation trusted clean stream exits as success + without any verification, which produced false-positive notifications + when SSE keep-alive timeouts closed the stream early. + + Race-safety note: ``runs.get`` returning ``"error"`` immediately after + a clean stream close can be a transient state for an actually-successful + run (server hasn't finalized the writeback). We trust the absence of + in-band error event over a stale ``runs.get="error"`` — see the + ``status == "error" and not saw_error_event`` branch below. """ - stream_failed = False - saw_error_event = False - try: - async for chunk in client.runs.join_stream( - thread_id=thread_id, run_id=run_id, stream_mode="values" - ): - ev = getattr(chunk, "event", None) - data = getattr(chunk, "data", None) - if ev == "error": - # Authoritative in-band error signal from langgraph dev. - saw_error_event = True - logger.info("Watcher saw error event for task %s: %r", thread_id, data) - except Exception: - stream_failed = True - logger.warning("Watcher stream failed for task %s", thread_id, exc_info=True) + for attempt in range(_MAX_RECONNECT_ATTEMPTS + 1): + stream_failed = False + saw_error_event = False + try: + async for chunk in client.runs.join_stream( + thread_id=thread_id, run_id=run_id, stream_mode="values" + ): + ev = getattr(chunk, "event", None) + data = getattr(chunk, "data", None) + if ev == "error": + saw_error_event = True + logger.info( + "Watcher saw error event for task %s: %r", thread_id, data + ) + except Exception: + stream_failed = True + logger.warning( + "Watcher stream failed for task %s", thread_id, exc_info=True + ) - if saw_error_event: - status = "error" - elif not stream_failed: - # Clean stream exit, no error event → trust the server-side success. - status = "success" - else: - # Stream errored without delivering a final state — fall back to - # runs.get (best-effort). REJECT non-terminal statuses (pending / - # running / unknown): the run is still alive, we shouldn't notify - # at all. Returning early without enqueueing prevents the - # "⚠ pending" notification we observed when the stream failed - # mid-flight. + if saw_error_event: + status = "error" + break + + # Verify with server before deciding the run is done — clean stream + # close does NOT guarantee terminal state. try: run = await client.runs.get(thread_id=thread_id, run_id=run_id) - raw_status = run.get("status", "") - if raw_status in TERMINAL_STATUSES: - status = raw_status - else: - logger.info( - "Watcher fallback got non-terminal status %r for task %s; " - "skipping notification (run still alive)", - raw_status, + raw = run.get("status", "") + except Exception: + # Cannot verify terminal state. Defaulting to "success" here would + # reintroduce the false-positive class this watcher exists to + # prevent (clean stream + transient runs.get failure → unverified + # success). Retry within the reconnect budget; on exhaustion drop + # the notification rather than guess. + if attempt >= _MAX_RECONNECT_ATTEMPTS: + logger.warning( + "Watcher runs.get failed for task %s after %d reconnects; " + "unable to verify terminal state, skipping notification", thread_id, + _MAX_RECONNECT_ATTEMPTS, + exc_info=True, ) return - except Exception: - status = "error" + logger.warning( + "Watcher runs.get failed for task %s; retrying after backoff " + "(attempt %d)", + thread_id, + attempt + 1, + exc_info=True, + ) + await asyncio.sleep(min(0.25 * (attempt + 1), 2.0)) + continue + + if raw not in TERMINAL_STATUSES: + # Non-terminal status — includes the documented ``pending`` / + # ``running`` values AND any future / unknown status the SDK may + # introduce. Stream closed early but run is not done; re-join + # unless we've exhausted attempts. Treating unknown statuses as + # non-terminal is the safe default — better to retry once more + # than to enqueue a false-positive on an unrecognized state. + if attempt >= _MAX_RECONNECT_ATTEMPTS: + logger.warning( + "Watcher gave up on task %s after %d reconnects " + "(server still reports %r); skipping notification", + thread_id, + _MAX_RECONNECT_ATTEMPTS, + raw, + ) + return + logger.info( + "Watcher SSE closed for task %s but run reports %r; " + "re-joining (attempt %d)", + thread_id, + raw, + attempt + 1, + ) + continue + + if raw == "error": + # Race-safe interpretation: no in-band error event → trust the + # absence over the server-side ``error`` (likely transient + # writeback state for a successful run). Stream-failure path + # is the one case where we DO trust ``error`` — the stream + # blowing up usually means something genuinely went wrong. + status = "error" if stream_failed else "success" + break + + # success / timeout / interrupted — trust authoritative terminal status. + status = raw + break + else: + # Loop exhausted without a break — should be unreachable because the + # re-join branch returns explicitly when attempts are exhausted, but + # guard against future refactors. + return notification = AsyncTaskNotification( task_id=thread_id, diff --git a/tests/test_async_notifier.py b/tests/test_async_notifier.py index c0f194e..d1f8e54 100644 --- a/tests/test_async_notifier.py +++ b/tests/test_async_notifier.py @@ -737,8 +737,37 @@ def test_watcher_reports_error_on_in_band_error_event(run_async): client.runs.get.assert_not_awaited() -def test_watcher_clean_exit_without_error_event_is_success(run_async): - """Clean stream exit with no error event → success (no runs.get poll).""" +def test_watcher_clean_exit_with_runs_get_success_is_success(run_async): + """Clean stream exit + runs.get reports success → status=success.""" + + async def fake_stream(*a, **kw): + yield SimpleNamespace( + event="values", data={"messages": [{"type": "ai", "content": "ok"}]} + ) + + client = MagicMock() + client.runs.join_stream = fake_stream + client.runs.get = AsyncMock(return_value={"status": "success"}) + + _drain_all(async_notifier) + run_async(async_notifier.watch_run_and_notify(client, "thrS", "rS", "agentS")) + + notif = async_notifier._notification_queue.get_nowait() + assert notif.status == "success" + client.runs.get.assert_awaited_once() + + +def test_watcher_clean_exit_with_runs_get_error_is_race_safe(run_async): + """Clean stream exit + no in-band error event + runs.get returns 'error' + → status=success (race-safe). + + Server-side state writeback can transiently report 'error' for an + actually-successful run between SSE close and final-state finalization. + The absence of an in-band error event is authoritative — the run did + not actually error. This test guards against re-introducing the race + we hit when an earlier 'always-poll runs.get' attempt blindly trusted + the runs.get value. + """ async def fake_stream(*a, **kw): yield SimpleNamespace( @@ -747,8 +776,6 @@ def test_watcher_clean_exit_without_error_event_is_success(run_async): client = MagicMock() client.runs.join_stream = fake_stream - # If we (wrongly) poll runs.get and it returned "error", the test would - # fail — proving we no longer have the timing race. client.runs.get = AsyncMock(return_value={"status": "error"}) _drain_all(async_notifier) @@ -756,7 +783,160 @@ def test_watcher_clean_exit_without_error_event_is_success(run_async): notif = async_notifier._notification_queue.get_nowait() assert notif.status == "success" - client.runs.get.assert_not_awaited() + + +def test_watcher_clean_exit_with_runs_get_running_drops_notification(run_async): + """Reproduces the production bug: clean SSE close while run is still + actually running (HTTP keep-alive timeout under concurrency). + + Pre-fix: watcher trusted clean stream exit as 'success' and enqueued a + false-positive notification for a still-running task. + + Post-fix: watcher verifies via runs.get and re-joins the stream until + either a terminal status arrives or the reconnect budget is exhausted. + With a mock that perpetually closes cleanly + reports 'running', the + watcher exhausts retries and enqueues nothing. + """ + + async def fake_stream(*a, **kw): + # SSE closes cleanly after one chunk — simulates HTTP keep-alive + # timeout where the server drops the long-poll without an error. + yield SimpleNamespace(event="values", data={"messages": []}) + + client = MagicMock() + client.runs.join_stream = fake_stream + client.runs.get = AsyncMock(return_value={"status": "running"}) + + _drain_all(async_notifier) + run_async( + async_notifier.watch_run_and_notify( + client, "thr-bug", "rB", "data-analysis-agent" + ) + ) + + # No notification should have been enqueued anywhere. + assert _drain_one_queue_helper(async_notifier._unrouted_queue) == [] + assert _drain_one_queue_helper(async_notifier._notification_queue) == [] + for q in async_notifier._notifications_by_thread.values(): + assert _drain_one_queue_helper(q) == [] + # runs.get must have been polled at least once (the verify step). + assert client.runs.get.await_count >= 1 + + +def test_watcher_unknown_status_treated_as_non_terminal(run_async): + """Future / unrecognized status values should trigger a re-join, not a + false-positive notification. + + If the SDK introduces a new non-terminal status (e.g. ``queued``, + ``scheduled``) the watcher must NOT silently default to ``success`` — + that would re-introduce the same class of bug we just fixed. The + safe-default policy: anything outside ``TERMINAL_STATUSES`` is treated + as ``running``-equivalent and triggers re-join. + """ + + async def fake_stream(*a, **kw): + yield SimpleNamespace(event="values", data={"messages": []}) + + client = MagicMock() + client.runs.join_stream = fake_stream + # First call: hypothetical future status. Second call: actual completion. + client.runs.get = AsyncMock( + side_effect=[{"status": "queued"}, {"status": "success"}] + ) + + _drain_all(async_notifier) + run_async(async_notifier.watch_run_and_notify(client, "thrU", "rU", "agentU")) + + notif = async_notifier._notification_queue.get_nowait() + assert notif.status == "success" + # Re-joined because the unknown status was not terminal. + assert client.runs.get.await_count == 2 + + +def test_watcher_runs_get_persistent_failure_drops_notification(run_async, monkeypatch): + """If ``runs.get`` keeps raising, the watcher cannot verify terminal + state and MUST drop the notification rather than default to + ``"success"`` — otherwise a transient server outage reintroduces the + same false-positive class this watcher exists to prevent.""" + + async def fake_stream(*a, **kw): + yield SimpleNamespace(event="values", data={"messages": []}) + + client = MagicMock() + client.runs.join_stream = fake_stream + client.runs.get = AsyncMock(side_effect=RuntimeError("server unreachable")) + + # Skip the backoff sleeps to keep this test fast. + async def _no_sleep(*a, **kw): + return None + + monkeypatch.setattr(async_notifier.asyncio, "sleep", _no_sleep) + + _drain_all(async_notifier) + run_async(async_notifier.watch_run_and_notify(client, "thrG", "rG", "agentG")) + + # No notification — watcher exhausted the reconnect budget. Check every + # queue routing could send to so a future routing change can't make this + # test silently false-pass. + assert _drain_one_queue_helper(async_notifier._unrouted_queue) == [] + assert _drain_one_queue_helper(async_notifier._notification_queue) == [] + if hasattr(async_notifier, "_notifications_by_thread"): + for q in async_notifier._notifications_by_thread.values(): + assert _drain_one_queue_helper(q) == [] + # 1 initial + _MAX_RECONNECT_ATTEMPTS retries = 11 calls total. + assert client.runs.get.await_count == async_notifier._MAX_RECONNECT_ATTEMPTS + 1 + + +def test_watcher_runs_get_transient_failure_recovers(run_async, monkeypatch): + """A single ``runs.get`` failure followed by a successful response on + retry must produce a correct notification — verifies the bounded + retry path actually recovers from transient outages instead of just + eating notifications.""" + + async def fake_stream(*a, **kw): + yield SimpleNamespace(event="values", data={"messages": []}) + + client = MagicMock() + client.runs.join_stream = fake_stream + # First call raises (transient), second call returns terminal status. + client.runs.get = AsyncMock( + side_effect=[RuntimeError("blip"), {"status": "success"}] + ) + + async def _no_sleep(*a, **kw): + return None + + monkeypatch.setattr(async_notifier.asyncio, "sleep", _no_sleep) + + _drain_all(async_notifier) + run_async(async_notifier.watch_run_and_notify(client, "thrT", "rT", "agentT")) + + notif = async_notifier._notification_queue.get_nowait() + assert notif.status == "success" + assert client.runs.get.await_count == 2 + + +def test_watcher_re_joins_stream_until_terminal_status(run_async): + """When runs.get returns 'running' on attempt N but a terminal status + on attempt N+1, the watcher re-joins, observes the terminal status, + and enqueues the notification correctly.""" + + async def fake_stream(*a, **kw): + yield SimpleNamespace(event="values", data={"messages": []}) + + client = MagicMock() + client.runs.join_stream = fake_stream + # First call: still running. Second call: success. + client.runs.get = AsyncMock( + side_effect=[{"status": "running"}, {"status": "success"}] + ) + + _drain_all(async_notifier) + run_async(async_notifier.watch_run_and_notify(client, "thrR", "rR", "agentR")) + + notif = async_notifier._notification_queue.get_nowait() + assert notif.status == "success" + assert client.runs.get.await_count == 2 # ============================================================================