Allow manual session release to restore node availability

backend/monitor.py

@@ -9,13 +9,20 @@ async def check_node_status(node_id: str, node: dict, redis_client):
         async with httpx.AsyncClient(timeout=3) as client:
             resp = await client.get(url)
             if resp.status_code == 200:
-                data = json.loads(await redis_client.hget("nodes", node_id))
-                data["status"] = "online"
+                node_json = await redis_client.hget("nodes", node_id)
+                if not node_json:
+                    return
+                data = json.loads(node_json)
+                if data.get("status") != "busy":
+                    data["status"] = "online"
                 await redis_client.hset("nodes", node_id, json.dumps(data))
             else:
                 raise Exception()
     except Exception:
-        data = json.loads(await redis_client.hget("nodes", node_id))
+        node_json = await redis_client.hget("nodes", node_id)
+        if not node_json:
+            return
+        data = json.loads(node_json)
         data["status"] = "offline"
         await redis_client.hset("nodes", node_id, json.dumps(data))
 

backend/node_manager.py

@@ -11,6 +11,15 @@
 async def register_node(node: Dict):
     node_id = node.get("id", str(uuid.uuid4()))
     node["id"] = node_id
+
+    # Ensure max_sessions and active_sessions are tracked numerically
+    max_sessions = int(node.get("max_sessions", 1))
+    active_sessions = int(node.get("active_sessions", 0))
+
+    node["max_sessions"] = max_sessions
+    node["active_sessions"] = active_sessions
+    node.setdefault("status", "online")
+
     await redis_client.hset("nodes", node_id, json.dumps(node))
     return {"message": "Node registered", "id": node_id}
 

backend/proxy_router.py

@@ -3,28 +3,105 @@
 import httpx
 import redis.asyncio as aioredis
 import json
+from typing import Optional
 
 router = APIRouter()
 redis_client = aioredis.from_url("redis://10.160.13.16:6379/0", decode_responses=True)
 
+SESSION_MAP_KEY = "session_map"
+
+
+def _extract_session_id_from_path(path: str) -> Optional[str]:
+    parts = [part for part in path.split("/") if part]
+    if len(parts) >= 2 and parts[0] == "session":
+        return parts[1]
+    return None
+
+
+def _extract_session_id_from_response(content: bytes) -> Optional[str]:
+    try:
+        payload = json.loads(content.decode())
+    except Exception:
+        return None
+
+    if not isinstance(payload, dict):
+        return None
+
+    if "sessionId" in payload:
+        return payload.get("sessionId")
+
+    value = payload.get("value")
+    if isinstance(value, dict):
+        return value.get("sessionId") or value.get("session_id")
+
+    return None
+
+
+async def _update_node_session_count(node_id: str, delta: int):
+    node_json = await redis_client.hget("nodes", node_id)
+    if not node_json:
+        return
+
+    node = json.loads(node_json)
+    active_sessions = int(node.get("active_sessions", 0)) + delta
+    if active_sessions < 0:
+        active_sessions = 0
+
+    node["active_sessions"] = active_sessions
+
+    max_sessions = int(node.get("max_sessions", 1))
+    if active_sessions >= max_sessions:
+        node["status"] = "busy"
+    elif node.get("status") == "busy":
+        node["status"] = "online"
+
+    await redis_client.hset("nodes", node_id, json.dumps(node))
+
+
+async def _cleanup_session(session_id: str, node_id: Optional[str]):
+    await redis_client.hdel(SESSION_MAP_KEY, session_id)
+    if node_id:
+        await _update_node_session_count(node_id, -1)
+
+
 async def forward_request(request: Request, path: str):
     body = await request.body()
     headers = dict(request.headers)
 
-    nodes = await redis_client.hgetall("nodes")
-    if not nodes:
-        raise HTTPException(status_code=503, detail="No nodes available")
+    session_id = _extract_session_id_from_path(path)
 
-    # Pick first available online node
     target_node = None
-    for node_data in nodes.values():
-        node = json.loads(node_data)
-        if node.get("status") == "online":
-            target_node = node
-            break
+    target_node_id = None
 
-    if not target_node:
-        raise HTTPException(status_code=503, detail="No online nodes available")
+    if session_id:
+        target_node_id = await redis_client.hget(SESSION_MAP_KEY, session_id)
+        if not target_node_id:
+            raise HTTPException(status_code=404, detail="Session not found")
+
+        node_json = await redis_client.hget("nodes", target_node_id)
+        if not node_json:
+            await redis_client.hdel(SESSION_MAP_KEY, session_id)
+            raise HTTPException(status_code=503, detail="Node unavailable for session")
+
+        target_node = json.loads(node_json)
+    else:
+        nodes = await redis_client.hgetall("nodes")
+        if not nodes:
+            raise HTTPException(status_code=503, detail="No nodes available")
+
+        for node_id, node_data in nodes.items():
+            node = json.loads(node_data)
+            status = node.get("status")
+            max_sessions = int(node.get("max_sessions", 1))
+            active_sessions = int(node.get("active_sessions", 0))
+
+            if status == "online" and active_sessions < max_sessions:
+                target_node = node
+                target_node_id = node_id
+                break
+
+        if not target_node:
+            raise HTTPException(status_code=503, detail="No nodes available for new session")
 
     target_url = f"http://{target_node['host']}:{target_node['port']}/wd/hub/{path}"
 
@@ -33,22 +110,58 @@ async def forward_request(request: Request, path: str):
             request.method, target_url, headers=headers, content=body
         )
 
-    return resp.content, resp.status_code, resp.headers
+    return resp.content, resp.status_code, resp.headers, target_node_id or target_node.get("id"), session_id
+
 
-# ✅ Add explicit route for session creation
 @router.api_route("/wd/hub/session", methods=["POST", "OPTIONS"])
 async def create_session(request: Request):
-    content, status, headers = await forward_request(request, "session")
+    content, status, headers, node_id, _ = await forward_request(request, "session")
+
+    if node_id and 200 <= status < 300:
+        session_id = _extract_session_id_from_response(content)
+        if session_id:
+            await redis_client.hset(SESSION_MAP_KEY, session_id, node_id)
+            await _update_node_session_count(node_id, 1)
+
     return Response(content=content, status_code=status, headers=dict(headers))
 
-# ✅ Add route for generic proxy handling (any other Appium path)
+
 @router.api_route("/wd/hub/{path:path}", methods=["GET", "POST", "DELETE", "PUT", "PATCH", "OPTIONS"])
 async def proxy_generic(request: Request, path: str):
-    content, status, headers = await forward_request(request, path)
+    content, status, headers, node_id, session_id = await forward_request(request, path)
+
+    if (
+        request.method == "DELETE"
+        and session_id
+        and (
+            200 <= status < 300
+            or status == 404
+            or status == 410
+        )
+    ):
+        await _cleanup_session(session_id, node_id)
+
     return Response(content=content, status_code=status, headers=dict(headers))
 
-# ✅ Add Selenium-style root route (no wd/hub prefix)
+
 @router.api_route("/session", methods=["POST", "OPTIONS"])
 async def selenium_create_session(request: Request):
-    content, status, headers = await forward_request(request, "session")
+    content, status, headers, node_id, _ = await forward_request(request, "session")
+
+    if node_id and 200 <= status < 300:
+        session_id = _extract_session_id_from_response(content)
+        if session_id:
+            await redis_client.hset(SESSION_MAP_KEY, session_id, node_id)
+            await _update_node_session_count(node_id, 1)
+
     return Response(content=content, status_code=status, headers=dict(headers))
+
+
+@router.delete("/sessions/{session_id}")
+async def release_session(session_id: str):
+    node_id = await redis_client.hget(SESSION_MAP_KEY, session_id)
+    if not node_id:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    await _cleanup_session(session_id, node_id)
+    return {"message": f"Session {session_id} terminated"}