feat: enhance agent identification and connection management across multiple modules

Author: Killea

src/db/crud.py

@@ -267,7 +267,11 @@ async def msg_post(
     author_id = None
     author_name = author
 
-    async with db.execute("SELECT id, name FROM agents WHERE id = ?", (author,)) as cur:
+    # Try to find agent by ID, name, or display_name
+    async with db.execute(
+        "SELECT id, name FROM agents WHERE id = ? OR name = ? OR display_name = ?",
+        (author, author, author)
+    ) as cur:
         row = await cur.fetchone()
         if row:
             actual_author = row["name"]

src/main.py

@@ -139,7 +139,7 @@ async def __call__(self, scope, receive, send):
 @app.get("/mcp/sse")
 async def mcp_sse_endpoint(request: Request):
     """MCP SSE endpoint consumed by MCP clients (Claude Desktop, Cursor, …)."""
-    from src.mcp_server import init_session_id
+    from src.mcp_server import init_session_id, clear_connection_agent
 
     # Initialize unique session ID for this SSE connection
     session_id = init_session_id()
@@ -161,6 +161,9 @@ async def mcp_sse_endpoint(request: Request):
         # Most are normal disconnects (anyio.ClosedResourceError, CancelledError…).
         # Log at DEBUG to avoid polluting the terminal.
         logger.debug("MCP SSE session ended: %s: %s", type(exc).__name__, exc)
+    finally:
+        # Clean up connection agent registry on disconnect
+        clear_connection_agent(session_id)
     return _SseCompletedResponse()
 
 

src/mcp_server.py

@@ -71,6 +71,12 @@ def get_connection_agent() -> tuple[str | None, str | None]:
         return agent_info["agent_id"], agent_info["token"]
     return None, None
 
+def clear_connection_agent(session_id: str) -> None:
+    """Clear agent identity for a session (call on SSE disconnect)."""
+    if session_id in _connection_agents:
+        agent_info = _connection_agents.pop(session_id)
+        logger.info(f"[clear_connection_agent] removed session {session_id[:8]}: agent_id={agent_info.get('agent_id')}")
+
 # Create the MCP server instance
 server = Server("AgentChatBus")
 

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

@@ -8,9 +8,16 @@
     }
 
     const secondsAgo = (now - activityTime) / 1000;
-    if (agent.last_activity === "msg_wait" && secondsAgo < 60) return "Waiting";
-    if (secondsAgo < 30) return "Active";
-    if (secondsAgo < 300) return "Idle";
+    
+    if (agent.last_activity === "msg_wait" && secondsAgo < 60) {
+      return "Waiting";
+    }
+    if (secondsAgo < 30) {
+      return "Active";
+    }
+    if (secondsAgo < 300) {
+      return "Idle";
+    }
     return agent.is_online ? "Idle" : "Offline";
   }
 

src/static/js/shared-agents.js

@@ -155,26 +155,44 @@
       const msgArea = document.getElementById("messages");
       if (msgArea) {
         const rows = msgArea.querySelectorAll("[data-author-id]");
+        // Use agent's real UUID as key for deduplication
+        const agentUuidMap = new Map();  // UUID -> agent object
+        
         rows.forEach((row) => {
           const authorId = row.getAttribute("data-author-id");
-          if (authorId && authorId !== "system" && authorId !== "human" && !participantIdMap.has(authorId)) {
-            const agent = allAgents.find((a) => a.id === authorId || a.agent_id === authorId);
+          if (authorId && authorId !== "system" && authorId !== "human") {
+            // Try to find agent by id (UUID), author_id, or name/display_name
+            // Also try partial match for name/display_name (e.g., "iFlow CLI" matches "iFlow CLI (glm-5)")
+            const agent = allAgents.find((a) => 
+              a.id === authorId || 
+              a.agent_id === authorId ||
+              a.name === authorId ||
+              a.display_name === authorId ||
+              a.author_id === authorId ||
+              (a.name && a.name.includes(authorId)) ||
+              (a.display_name && a.display_name.includes(authorId))
+            );
             if (agent) {
-              participantIdMap.set(authorId, agent);
+              // Dedupe by agent's real UUID
+              if (!agentUuidMap.has(agent.id)) {
+                agentUuidMap.set(agent.id, agent);
+              }
             } else {
-              participantIdMap.set(authorId, {
-                id: authorId,
-                display_name: authorId,
-                name: authorId,
-                is_online: false,
-              });
+              // Fallback: use authorId as pseudo-UUID
+              if (!agentUuidMap.has(authorId)) {
+                agentUuidMap.set(authorId, {
+                  id: authorId,
+                  display_name: authorId,
+                  name: authorId,
+                  is_online: false,
+                });
+              }
             }
           }
         });
+        // When a thread is selected, only show participants for that thread
+        participants = Array.from(agentUuidMap.values());
       }
-
-      // When a thread is selected, only show participants for that thread
-      participants = Array.from(participantIdMap.values());
     } else {
       // When no thread is selected, show only online agents
       participants = allAgents.filter((a) => a.is_online);

src/static/js/shared-api.js

@@ -9,7 +9,7 @@
       if (!response.ok) {
         console.warn(`[API] ${options.method || 'GET'} ${path} → HTTP ${response.status}`);
       }
-      return response.json();
+      return await response.json();
     } catch (err) {
       console.warn(`[API] ${options.method || 'GET'} ${path} → network error:`, err);
       return null;

src/tools/dispatch.py

@@ -290,22 +290,45 @@ async def handle_msg_wait(db, arguments: dict[str, Any]) -> list[types.Content]:
 
     logger.info(f"[msg_wait] explicit: agent_id={explicit_agent_id}, connection: agent_id={connection_agent_id}, final_agent_id={agent_id}")
 
-    if agent_id and token:
-        try:
-            result = await crud.agent_msg_wait(db, agent_id, token)
-            logger.info(f"[msg_wait] activity recorded: agent_id={agent_id}, result={result}")
-        except Exception as e:
-            logger.warning(f"[msg_wait] Failed to record activity for {agent_id}: {e}")
-    else:
-        logger.warning(f"[msg_wait] No credentials available: agent_id={agent_id}, token={'***' if token else None}")
+    # Heartbeat interval: refresh every 20 seconds to stay online (AGENT_HEARTBEAT_TIMEOUT is 30s)
+    HEARTBEAT_INTERVAL = 20.0
+
+    async def _refresh_heartbeat():
+        """Refresh heartbeat to keep agent online during long wait."""
+        if agent_id and token:
+            try:
+                await crud.agent_msg_wait(db, agent_id, token)
+                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}")
 
     async def _poll():
+        """Poll for new messages, refreshing heartbeat periodically."""
+        last_heartbeat = asyncio.get_event_loop().time()
         while True:
+            # Check for new messages
             msgs = await crud.msg_list(db, thread_id, after_seq=after_seq, include_system_prompt=False)
             if msgs:
                 return msgs
+            
+            # Refresh heartbeat if interval elapsed
+            now = asyncio.get_event_loop().time()
+            if now - last_heartbeat >= HEARTBEAT_INTERVAL:
+                await _refresh_heartbeat()
+                last_heartbeat = now
+            
             await asyncio.sleep(0.5)
 
+    # Record initial activity
+    if agent_id and token:
+        try:
+            result = await crud.agent_msg_wait(db, agent_id, token)
+            logger.info(f"[msg_wait] initial activity recorded: agent_id={agent_id}, result={result}")
+        except Exception as e:
+            logger.warning(f"[msg_wait] Failed to record activity for {agent_id}: {e}")
+    else:
+        logger.warning(f"[msg_wait] No credentials available: agent_id={agent_id}, token={'***' if token else None}")
+
     try:
         msgs = await asyncio.wait_for(_poll(), timeout=timeout_s)
     except asyncio.TimeoutError: