feat: Add batch heartbeat submission support (#106)

beacon_skill/relay.py

@@ -329,6 +329,70 @@ def heartbeat(
             },
         }
 
+
+    def batch_heartbeat(
+        self,
+        heartbeats: List[Dict[str, Any]],
+        default_token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Process a batched relay heartbeat.
+
+        Args:
+            heartbeats: List of dicts, each with 'agent_id', 'status', and optionally 'token', 'health'.
+            default_token: Token to use if a heartbeat doesn't provide one.
+        """
+        results = {}
+        agents = self._load_agents()
+        now = int(time.time())
+        dirty = False
+
+        for hb in heartbeats:
+            agent_id = hb.get("agent_id")
+            token = hb.get("token") or default_token
+            status = hb.get("status", "alive")
+            health = hb.get("health")
+
+            if not agent_id:
+                continue
+
+            agent_data = agents.get(agent_id)
+            if not agent_data:
+                results[agent_id] = {"error": "Agent not registered", "code": "NOT_FOUND"}
+                continue
+
+            if agent_data.get("relay_token") != token:
+                results[agent_id] = {"error": "Invalid relay token", "code": "AUTH_FAILED"}
+                continue
+
+            if agent_data.get("token_expires", 0) < now:
+                results[agent_id] = {"error": "Token expired — re-register", "code": "TOKEN_EXPIRED"}
+                continue
+
+            # Update heartbeat state
+            agent_data["last_heartbeat"] = now
+            agent_data["beat_count"] = agent_data.get("beat_count", 0) + 1
+            agent_data["status"] = status
+
+            # Refresh token TTL
+            agent_data["token_expires"] = now + RELAY_TOKEN_TTL_S
+
+            if health:
+                agent_data.setdefault("metadata", {})
+                agent_data["metadata"]["last_health"] = health
+
+            results[agent_id] = {
+                "ok": True, 
+                "expires_in": RELAY_TOKEN_TTL_S,
+                "agent_id": agent_id,
+                "beat_count": agent_data["beat_count"]
+            }
+            dirty = True
+
+        if dirty:
+            self._save_agents(agents)
+
+        return {"batch_results": results}
+
     # ── Discovery ──
 
     def discover(self, provider: Optional[str] = None, capability: Optional[str] = None) -> List[Dict[str, Any]]:

beacon_skill/transports/relay.py

@@ -184,6 +184,28 @@ def heartbeat_seo(
         )
         return resp.json()
 
+    def batch_heartbeat(
+        self,
+        heartbeats: List[Dict[str, Any]],
+        token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Send multiple heartbeats in a single request.
+
+        Args:
+            heartbeats: List of dicts, each with 'agent_id', 'status', and optionally 'health', 'token'.
+            token: Global auth token for the batch.
+        """
+        body = {
+            "heartbeats": heartbeats
+        }
+        resp = requests.post(
+            self._url("/relay/heartbeat/batch"),
+            json=body,
+            headers=self._headers(token or ""),
+            timeout=self.timeout_s,
+        )
+        return resp.json()
+
     def discover(
         self,
         provider: Optional[str] = None,

tests/test_relay_batch_heartbeat.py

@@ -0,0 +1,47 @@
+import pytest
+import time
+from beacon_skill.relay import RelayManager
+from beacon_skill.transports.relay import RelayClient
+
+def test_relay_batch_heartbeat(tmp_path):
+    mgr = RelayManager(data_dir=tmp_path)
+
+    now = int(time.time())
+    # Register a few agents directly into the manager
+    agents = {
+        "agent_1": {
+            "relay_token": "token1",
+            "token_expires": now + 86400,
+            "status": "dormant"
+        },
+        "agent_2": {
+            "relay_token": "token2",
+            "token_expires": now + 86400,
+            "status": "dormant"
+        }
+    }
+    mgr._save_agents(agents)
+
+    heartbeats = [
+        {"agent_id": "agent_1", "status": "active", "token": "token1"},
+        {"agent_id": "agent_2", "status": "idle", "token": "token2"},
+        {"agent_id": "agent_3", "status": "busy", "token": "wrong"} # not registered
+    ]
+
+    # We patch RELAY_TOKEN_TTL_S to be accessible in globals if needed, or it's built-in
+
+    result = mgr.batch_heartbeat(heartbeats)
+
+    assert "batch_results" in result
+    assert "agent_1" in result["batch_results"]
+    assert "agent_2" in result["batch_results"]
+    assert "agent_3" in result["batch_results"]
+
+    assert result["batch_results"]["agent_1"].get("ok") is True
+    assert result["batch_results"]["agent_2"].get("ok") is True
+    assert result["batch_results"]["agent_3"].get("code") == "NOT_FOUND"
+
+    updated = mgr._load_agents()
+    assert updated["agent_1"]["status"] == "active"
+    assert updated["agent_2"]["status"] == "idle"
+