feat: Update SSE session management and agent status handling for improved connection accuracy

Author: Killea

src/mcp_server.py

@@ -58,9 +58,12 @@ def init_session_id() -> str:
 
 # Live SSE session registry (in-memory only, resets on server restart).
 # Contains session_ids that currently have an active SSE TCP connection.
-# SSE agents in this set are provably online regardless of heartbeat timeout.
+# Sessions are refreshed on each tool call POST and during msg_wait heartbeats.
+# After the last refresh, sessions are considered stale after _SSE_STALE_SECONDS.
+# Kept deliberately short (60s) so that agents go offline quickly after msg_wait
+# ends or they stop making tool calls, rather than staying "live" for 5+ minutes.
 _active_sse_sessions: dict[str, float] = {}
-_SSE_STALE_SECONDS = max(120, MSG_WAIT_TIMEOUT + 15)
+_SSE_STALE_SECONDS = 60
 def get_session_id() -> str | None:
     """Get session ID for this SSE connection."""
     return _session_id.get()

src/static/js/shared-agent-status.js

@@ -3,18 +3,29 @@
    * Determine agent state using SSE connection presence as the primary signal.
    *
    * States (priority order):
-   *   Listening — SSE connected + last_activity is msg_wait   (green ⏳)
-   *   Working   — SSE connected + NOT in msg_wait             (amber ⚡)
-   *   Offline   — no SSE + heartbeat expired                  (dark ⚫)
+   *   Listening — SSE connected + last_activity is msg_wait              (⏳)
+   *   Working   — SSE connected + NOT msg_wait + activity within 60s     (⚡)
+   *   Idle      — SSE connected + NOT msg_wait + activity older than 60s (🔌)
+   *               (SSE open but agent stopped responding — may be processing
+   *               a long task or is a stale connection)
+   *   Offline   — no SSE + heartbeat expired                             (⚫)
    *
    * For stdio agents (no SSE): fall back to heartbeat-based detection.
    * stdio agents will show Listening when is_online=true, Offline otherwise.
    */
+  const WORKING_TIMEOUT_S = 60;
+
   function getAgentState(agent) {
     if (!agent) return "Offline";
 
     if (agent.is_sse_connected) {
       if (agent.last_activity === "msg_wait") return "Listening";
+      // Check how long ago the last activity was
+      const lastActivityTime = agent.last_activity_time ? new Date(agent.last_activity_time) : null;
+      if (lastActivityTime) {
+        const elapsedS = (Date.now() - lastActivityTime.getTime()) / 1000;
+        if (elapsedS > WORKING_TIMEOUT_S) return "Idle";
+      }
       return "Working";
     }
 
@@ -25,10 +36,10 @@
 
   /**
    * Return the state emoji for display in the card.
-   * ⏳ Listening, ⚡ Working, ⚫ Offline
+   * ⏳ Listening, ⚡ Working, 🔌 Idle, ⚫ Offline
    */
   function getStateEmoji(state) {
-    const map = { Listening: "⏳", Working: "⚡", Offline: "⚫" };
+    const map = { Listening: "⏳", Working: "⚡", Idle: "🔌", Offline: "⚫" };
     return map[state] || "⚫";
   }
 

src/tools/dispatch.py

@@ -816,6 +816,15 @@ async def _refresh_heartbeat() -> None:
                 logger.debug(f"[msg_wait] heartbeat refreshed for agent_id={agent_id}")
             except Exception as e:
                 logger.warning(f"[msg_wait] Failed to refresh heartbeat for {agent_id}: {e}")
+        # Also refresh the in-process SSE session timestamp so is_agent_sse_connected()
+        # stays true during active msg_wait polling. Without this, the short
+        # _SSE_STALE_SECONDS window would expire mid-wait and flip the agent offline.
+        try:
+            session_id = src.mcp_server.get_session_id()
+            if session_id:
+                src.mcp_server.mark_sse_connected(session_id)
+        except Exception:
+            pass
 
     if agent_id and token:
         try: