ouyangbo/EvoScientist-Multi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: Improve watcher logic to prevent false-positive notifications on… (#216)

Browse Source 
* fix: Improve watcher logic to prevent false-positive notifications on clean stream exits

* fix: Update watcher logic to drop notifications on persistent runs.get failures

* fix: Refactor test for watcher persistent failure notification handling

* fix: Enhance watcher test to validate all notification queues are empty after reconnect budget exhaustion
This commit is contained in:
Xi Zhang
2026-05-07 23:16:36 +02:00
committed by GitHub
parent 80f1f4fa0f
commit 4e04ac5b72
2 changed files with 303 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
EvoScientist/cli/async_notifier.py
+118 -48
View File
@@ -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,
tests/test_async_notifier.py
+185 -5
View File
@@ -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
# ============================================================================