Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions api/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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".
Expand Down
41 changes: 41 additions & 0 deletions api/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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}
Expand Down
14 changes: 14 additions & 0 deletions api/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 = (
Expand Down Expand Up @@ -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:
Expand All @@ -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`
Expand Down
50 changes: 50 additions & 0 deletions tests/test_run_lifecycle_health.py
Original file line number Diff line number Diff line change
@@ -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)
Loading