diff --git a/api/config.py b/api/config.py index 5a592dea229..d2675e69eb7 100644 --- a/api/config.py +++ b/api/config.py @@ -3681,8 +3681,48 @@ def create_stream_channel() -> StreamChannel: STREAM_LIVE_TOOL_CALLS: dict = {} # stream_id -> live tool calls accumulated during streaming (#1361 §B) STREAM_GOAL_RELATED: dict = {} # stream_id -> bool: only evaluate goal for goal-related turns (#1932) PENDING_GOAL_CONTINUATION: set = set() # session_ids awaiting a goal continuation turn (#1932) + +# Active agent-run registry. This intentionally tracks worker lifecycle rather +# than SSE lifecycle: cancel/reconnect may remove STREAMS while the worker is +# still unwinding, blocked in a provider call, or waiting for delegated work. +ACTIVE_RUNS: dict = {} +ACTIVE_RUNS_LOCK = threading.Lock() +LAST_RUN_FINISHED_AT: float | None = None SERVER_START_TIME = time.time() + +def register_active_run(stream_id: str, **metadata) -> None: + """Mark a WebUI agent worker as alive until its outer finally exits.""" + if not stream_id: + return + now = time.time() + entry = dict(metadata or {}) + entry.setdefault("stream_id", stream_id) + entry.setdefault("started_at", now) + entry.setdefault("phase", "running") + with ACTIVE_RUNS_LOCK: + ACTIVE_RUNS[stream_id] = entry + + +def update_active_run(stream_id: str, **metadata) -> None: + """Update active-run metadata without creating a new run implicitly.""" + if not stream_id: + return + with ACTIVE_RUNS_LOCK: + entry = ACTIVE_RUNS.get(stream_id) + if entry is not None: + entry.update(metadata) + + +def unregister_active_run(stream_id: str) -> None: + """Remove a worker from the active-run registry and record idle start.""" + if not stream_id: + return + global LAST_RUN_FINISHED_AT + with ACTIVE_RUNS_LOCK: + ACTIVE_RUNS.pop(stream_id, None) + LAST_RUN_FINISHED_AT = time.time() + # Agent cache: reuse AIAgent across messages in the same WebUI session so that # _user_turn_count survives between turns. This mirrors the gateway's # _agent_cache pattern and is required for injectionFrequency: "first-turn". diff --git a/api/routes.py b/api/routes.py index 33567d0dbb3..926250fcef6 100644 --- a/api/routes.py +++ b/api/routes.py @@ -2529,6 +2529,39 @@ def _streams_lock_health(timeout_seconds: float = 0.5) -> dict: STREAMS_LOCK.release() +def _run_lifecycle_health() -> dict: + """Return active worker-run state independent of SSE stream presence.""" + # Import the module rather than relying only on imported scalar aliases so + # LAST_RUN_FINISHED_AT stays fresh after unregister_active_run() updates it. + from api import config as _live_config + + now = time.time() + with _live_config.ACTIVE_RUNS_LOCK: + runs = [] + for stream_id, raw in (_live_config.ACTIVE_RUNS or {}).items(): + item = dict(raw or {}) + started_at = item.get("started_at") + try: + age = max(0.0, now - float(started_at)) + except Exception: + age = 0.0 + item.setdefault("stream_id", stream_id) + item["age_seconds"] = round(age, 1) + runs.append(item) + last_finished = _live_config.LAST_RUN_FINISHED_AT + runs.sort(key=lambda item: float(item.get("started_at") or 0.0)) + payload = { + "active_runs": len(runs), + "runs": runs, + "last_run_finished_at": last_finished, + } + if runs: + payload["oldest_run_age_seconds"] = runs[0].get("age_seconds", 0.0) + elif last_finished: + payload["idle_seconds_since_last_run"] = round(max(0.0, now - float(last_finished)), 1) + return payload + + def _deep_health_checks(stream_check: dict | None = None) -> tuple[dict, bool]: """Run cheap probes that exercise the state paths used by the UI shell. @@ -2609,13 +2642,21 @@ def _deep_health_checks(stream_check: dict | None = None) -> tuple[dict, bool]: def _handle_health(handler, parsed): deep = parse_qs(parsed.query or "").get("deep", [""])[0].lower() in {"1", "true", "yes", "on"} stream_check = _streams_lock_health() + run_check = _run_lifecycle_health() payload = { "status": "ok" if stream_check.get("status") == "ok" else "degraded", "sessions": len(SESSIONS), "active_streams": int(stream_check.get("active_streams") or 0), + "active_runs": int(run_check.get("active_runs") or 0), + "runs": run_check.get("runs", []), + "last_run_finished_at": run_check.get("last_run_finished_at"), "uptime_seconds": round(time.time() - SERVER_START_TIME, 1), "accept_loop": _accept_loop_health(handler), } + if "oldest_run_age_seconds" in run_check: + payload["oldest_run_age_seconds"] = run_check["oldest_run_age_seconds"] + if "idle_seconds_since_last_run" in run_check: + payload["idle_seconds_since_last_run"] = run_check["idle_seconds_since_last_run"] if deep: if stream_check.get("status") != "ok": payload["checks"] = {"streams_lock": stream_check} diff --git a/api/streaming.py b/api/streaming.py index 7968c6edbf6..8e231041a69 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -26,6 +26,7 @@ STREAM_GOAL_RELATED, PENDING_GOAL_CONTINUATION, LOCK, SESSIONS, SESSION_DIR, _get_session_agent_lock, _set_thread_env, _clear_thread_env, + register_active_run, update_active_run, unregister_active_run, SESSION_AGENT_LOCKS, SESSION_AGENT_LOCKS_LOCK, resolve_model_provider, resolve_custom_provider_connection, @@ -2006,6 +2007,16 @@ def _run_agent_streaming( q = STREAMS.get(stream_id) if q is None: return + register_active_run( + stream_id, + session_id=session_id, + started_at=time.time(), + phase="starting", + workspace=str(workspace), + model=model, + provider=model_provider, + ephemeral=bool(ephemeral), + ) s = None _rt = {} old_cwd = None @@ -2187,6 +2198,7 @@ def _agent_status_callback(kind, message): _agent_lock = None try: s = get_session(session_id) + update_active_run(stream_id, phase="running", session_id=session_id) s.workspace = str(Path(workspace).expanduser().resolve()) s.model = model provider_context = ( @@ -3882,6 +3894,7 @@ def _periodic_checkpoint(): if (s is not None and getattr(s, 'active_stream_id', None) == stream_id and getattr(s, 'pending_user_message', None)): + update_active_run(stream_id, phase="finalizing") _last_resort_sync_from_core(s, stream_id, _agent_lock) _clear_thread_env() # TD1: always clear thread-local context with STREAMS_LOCK: @@ -3892,6 +3905,7 @@ def _periodic_checkpoint(): STREAM_REASONING_TEXT.pop(stream_id, None) # Clean up reasoning trace (#1361 §A) STREAM_LIVE_TOOL_CALLS.pop(stream_id, None) # Clean up tool calls (#1361 §B) STREAM_GOAL_RELATED.pop(stream_id, None) # Clean up goal-related flag (#1932) + unregister_active_run(stream_id) # NOTE: do NOT discard PENDING_GOAL_CONTINUATION here. The marker # is set by goal_continue (line ~3328) inside the SAME function # call and consumed atomically by `_start_chat_stream_for_session` diff --git a/tests/test_run_lifecycle_health.py b/tests/test_run_lifecycle_health.py new file mode 100644 index 00000000000..8913ade7f4d --- /dev/null +++ b/tests/test_run_lifecycle_health.py @@ -0,0 +1,50 @@ +"""Regression coverage for restart-safety run lifecycle reporting.""" + +import time + + +def test_health_counts_active_runs_even_when_no_sse_streams(): + """A worker run can outlive its SSE channel; health must expose the run.""" + from api import config, routes + + with config.STREAMS_LOCK: + config.STREAMS.clear() + with config.ACTIVE_RUNS_LOCK: + config.ACTIVE_RUNS.clear() + config.ACTIVE_RUNS["stream-1"] = { + "stream_id": "stream-1", + "session_id": "session-1", + "started_at": time.time() - 42, + "phase": "running", + } + + try: + stream_check = routes._streams_lock_health() + run_check = routes._run_lifecycle_health() + + assert stream_check["active_streams"] == 0 + assert run_check["active_runs"] == 1 + assert run_check["oldest_run_age_seconds"] >= 40 + assert run_check["runs"][0]["session_id"] == "session-1" + finally: + with config.ACTIVE_RUNS_LOCK: + config.ACTIVE_RUNS.clear() + + +def test_run_registry_unregister_records_last_finished_time(): + """Guards need a grace window after the last real worker exits.""" + from api import config + + with config.ACTIVE_RUNS_LOCK: + config.ACTIVE_RUNS.clear() + config.LAST_RUN_FINISHED_AT = None + + config.register_active_run("stream-2", session_id="session-2", phase="starting") + with config.ACTIVE_RUNS_LOCK: + assert "stream-2" in config.ACTIVE_RUNS + + config.unregister_active_run("stream-2") + + with config.ACTIVE_RUNS_LOCK: + assert "stream-2" not in config.ACTIVE_RUNS + assert isinstance(config.LAST_RUN_FINISHED_AT, float)