diff --git a/api/routes.py b/api/routes.py index 64c5c6d9f..122ea3137 100644 --- a/api/routes.py +++ b/api/routes.py @@ -14139,8 +14139,9 @@ def _active_run_stream_for_session(session_id: str | None) -> str | None: now = time.time() try: from api import config as _live_config + stale_stream_ids = [] with _live_config.ACTIVE_RUNS_LOCK: - for run_stream_id, raw in (_live_config.ACTIVE_RUNS or {}).items(): + for run_stream_id, raw in list((_live_config.ACTIVE_RUNS or {}).items()): stream_id = str((raw or {}).get("stream_id") or run_stream_id or "").strip() run_sid = str((raw or {}).get("session_id") or "").strip() if run_sid != sid or not stream_id: @@ -14149,11 +14150,14 @@ def _active_run_stream_for_session(session_id: str | None) -> str | None: started_at = float((raw or {}).get("started_at") or 0) except (TypeError, ValueError): started_at = 0.0 - # Ignore a stale/wedged entry past the unwind ceiling so it can't - # block the session permanently. + # Reconcile stale/wedged entries past the unwind ceiling so they + # cannot keep health/recovery polling in a half-alive state. if started_at and (now - started_at) > ceiling: + stale_stream_ids.append(stream_id) continue return stream_id + for stale_stream_id in stale_stream_ids: + (_live_config.ACTIVE_RUNS or {}).pop(stale_stream_id, None) except Exception: return None return None diff --git a/api/streaming.py b/api/streaming.py index 7a2486bf6..adddf30ab 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -8903,6 +8903,11 @@ def cancel_stream(stream_id: str) -> bool: if active_run_entry and not active_run_session_id: active_run_session_id = str(active_run_entry.get("session_id") or "").strip() or None + # Mark the worker lifecycle registry immediately. The SSE maps may be popped + # below while the worker is still unwinding; ACTIVE_RUNS is what recovery / + # health polling sees during that detached window. + update_active_run(stream_id, phase="cancelling") + # Set WebUI layer cancel flag. Prefer the snapshot captured under the lock; # fall back to a fresh lookup for the ACTIVE_RUNS-only path (stream absent). flag = _snap_flag if _snap_flag is not None else cancel_flags.get(stream_id) diff --git a/tests/test_cancel_interrupt.py b/tests/test_cancel_interrupt.py index c2d35aac1..375d0ef52 100644 --- a/tests/test_cancel_interrupt.py +++ b/tests/test_cancel_interrupt.py @@ -127,6 +127,7 @@ class TestCancelInterrupt: result = cancel_stream(stream_id) assert result is True + assert ACTIVE_RUNS[stream_id]["phase"] == "cancelling" mock_agent.interrupt.assert_called_once_with("Cancelled by user") assert mock_session.active_stream_id is None assert mock_session.pending_user_message is None diff --git a/tests/test_stale_stream_cleanup.py b/tests/test_stale_stream_cleanup.py index 8d5756a96..2d9e929fe 100644 --- a/tests/test_stale_stream_cleanup.py +++ b/tests/test_stale_stream_cleanup.py @@ -319,6 +319,9 @@ def test_chat_start_not_permanently_blocked_by_stale_active_run(monkeypatch, tmp # The bounded guard should treat it as stale and NOT report it as blocking. assert routes._active_run_stream_for_session(session.session_id) is None + # It must also reconcile the zombie registry entry immediately so health / + # recovery polling does not keep advertising a half-alive run forever. + assert stale_stream_id not in config.ACTIVE_RUNS class NoopThread: def __init__(self, *args, **kwargs):