kairo-code 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

kairo/backend/tests/e2e/agents/test_agent_heartbeat.py

@@ -0,0 +1,520 @@
+"""
+End-to-end tests for Agent Heartbeat functionality.
+
+Tests cover:
+- Successful heartbeat updates last_heartbeat_at
+- Heartbeat with metrics records data
+- Heartbeat returns pending commands
+- State transitions (online, busy, idle, error)
+- Stale agent detection (mark offline after timeout)
+"""
+
+import pytest
+import uuid
+from datetime import datetime, timedelta, UTC
+from httpx import AsyncClient
+from sqlalchemy import select
+
+from backend.models.agent import Agent, AgentMetrics1m
+
+
+class TestHeartbeatBasic:
+    """Basic heartbeat functionality tests."""
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_success(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should successfully process heartbeat and return acknowledgment."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["acknowledged"] is True
+        assert "server_time" in data
+        assert "commands" in data
+        assert isinstance(data["commands"], list)
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_updates_last_heartbeat_at(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should update last_heartbeat_at timestamp on heartbeat."""
+        original_heartbeat = test_agent.last_heartbeat_at
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.last_heartbeat_at is not None
+        if original_heartbeat:
+            assert test_agent.last_heartbeat_at > original_heartbeat
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_updates_sdk_version(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should update SDK version on heartbeat."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+            "sdk_version": "2.5.0",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.sdk_version == "2.5.0"
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_updates_host_info(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should update host info on heartbeat."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+            "host_info": {
+                "hostname": "prod-server-01",
+                "ip": "10.0.1.50",
+                "os": "debian-11",
+                "memory_mb": 32768,
+                "memory_used_percent": 45.5,
+            },
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.host_info is not None
+        assert test_agent.host_info["hostname"] == "prod-server-01"
+        assert test_agent.host_info["memory_used_percent"] == 45.5
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_agent_not_found(
+        self, client: AsyncClient, api_key_headers: dict
+    ):
+        """Should return 404 for non-existent agent."""
+        fake_id = str(uuid.uuid4())
+        payload = {
+            "agent_id": fake_id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 404
+
+
+class TestHeartbeatStateTransitions:
+    """Tests for agent state transitions via heartbeat."""
+
+    @pytest.mark.asyncio
+    async def test_transition_to_online(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should transition agent to 'online' state."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.state == "online"
+        assert test_agent.last_online_at is not None
+
+    @pytest.mark.asyncio
+    async def test_transition_to_busy(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should transition agent to 'busy' state."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "busy",
+            "active_request": True,
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.state == "busy"
+
+    @pytest.mark.asyncio
+    async def test_transition_to_idle(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should transition agent to 'idle' state."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "idle",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.state == "idle"
+
+    @pytest.mark.asyncio
+    async def test_transition_to_error(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should transition agent to 'error' state with error details."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "error",
+            "last_error": {
+                "message": "Connection to LLM service failed",
+                "code": "LLM_CONNECTION_ERROR",
+            },
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.state == "error"
+        assert test_agent.last_error_at is not None
+        assert test_agent.last_error_message == "Connection to LLM service failed"
+
+    @pytest.mark.asyncio
+    async def test_first_connection_sets_timestamp(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should set first_connected_at on first heartbeat."""
+        assert test_agent.first_connected_at is None
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.first_connected_at is not None
+
+
+class TestHeartbeatMetrics:
+    """Tests for metrics recording via heartbeat."""
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_records_metrics(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should record metrics from heartbeat payload."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+            "metrics_since_last_heartbeat": {
+                "requests_completed": 10,
+                "requests_failed": 2,
+                "total_latency_ms": 5000,
+                "input_tokens": 10000,
+                "output_tokens": 5000,
+                "tool_calls": 15,
+                "tool_errors": 1,
+            },
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        # Check metrics were recorded
+        stmt = select(AgentMetrics1m).where(AgentMetrics1m.agent_id == test_agent.id)
+        result = await db_session.execute(stmt)
+        metrics = result.scalar_one_or_none()
+
+        assert metrics is not None
+        assert metrics.request_count == 12  # completed + failed
+        assert metrics.error_count == 2
+        assert metrics.input_tokens == 10000
+        assert metrics.output_tokens == 5000
+        assert metrics.tool_calls == 15
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_aggregates_metrics_same_bucket(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should aggregate metrics when multiple heartbeats occur in same minute."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+            "metrics_since_last_heartbeat": {
+                "requests_completed": 5,
+                "requests_failed": 1,
+                "input_tokens": 1000,
+                "output_tokens": 500,
+            },
+        }
+
+        # Send two heartbeats
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        # Check aggregated metrics
+        stmt = select(AgentMetrics1m).where(AgentMetrics1m.agent_id == test_agent.id)
+        result = await db_session.execute(stmt)
+        metrics = result.scalar_one_or_none()
+
+        assert metrics is not None
+        assert metrics.request_count == 12  # (5+1) * 2
+        assert metrics.error_count == 2  # 1 * 2
+        assert metrics.input_tokens == 2000  # 1000 * 2
+        assert metrics.output_tokens == 1000  # 500 * 2
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_without_metrics(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should handle heartbeat without metrics payload."""
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 200
+
+        # No metrics should be recorded
+        stmt = select(AgentMetrics1m).where(AgentMetrics1m.agent_id == test_agent.id)
+        result = await db_session.execute(stmt)
+        metrics = result.scalar_one_or_none()
+
+        assert metrics is None
+
+
+class TestHeartbeatCommands:
+    """Tests for pending command delivery via heartbeat."""
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_returns_pending_commands(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, agent_factory
+    ):
+        """Should return pending commands in heartbeat response."""
+        # Create a pending command
+        await agent_factory.create_command(
+            test_agent.id, command_type="restart", status="pending"
+        )
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert len(data["commands"]) == 1
+        assert data["commands"][0]["type"] == "restart"
+        assert "command_id" in data["commands"][0]
+        assert "signature" in data["commands"][0]
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_marks_commands_dispatched(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should mark commands as dispatched after delivery."""
+        command = await agent_factory.create_command(
+            test_agent.id, command_type="stop", status="pending"
+        )
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        await db_session.refresh(command)
+
+        assert command.status == "dispatched"
+        assert command.dispatched_at is not None
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_excludes_expired_commands(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, agent_factory
+    ):
+        """Should not return expired commands."""
+        # Create an expired command
+        await agent_factory.create_command(
+            test_agent.id,
+            command_type="restart",
+            status="pending",
+            expires_at=datetime.now(UTC) - timedelta(hours=1),
+        )
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert len(data["commands"]) == 0
+
+    @pytest.mark.asyncio
+    async def test_heartbeat_excludes_already_dispatched(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent, agent_factory
+    ):
+        """Should not return already dispatched commands."""
+        # Create a dispatched command
+        await agent_factory.create_command(
+            test_agent.id,
+            command_type="restart",
+            status="dispatched",
+        )
+
+        payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat", json=payload, headers=api_key_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert len(data["commands"]) == 0
+
+
+class TestStaleAgentDetection:
+    """Tests for stale agent detection and offline marking."""
+
+    @pytest.mark.asyncio
+    async def test_mark_stale_agent_offline(self, db_session, test_user):
+        """Should mark agents as offline when heartbeat times out."""
+        from backend.services.agent_service import AgentService
+
+        # Create an agent with old heartbeat
+        stale_time = datetime.now(UTC) - timedelta(minutes=5)
+        agent = Agent(
+            id=str(uuid.uuid4()),
+            user_id=test_user.id,
+            name="Stale Agent",
+            state="online",
+            status="online",
+            last_heartbeat_at=stale_time,
+        )
+        db_session.add(agent)
+        await db_session.commit()
+
+        # Run staleness check with 2-minute threshold
+        svc = AgentService(db_session)
+        count = await svc.mark_stale_agents_offline(threshold_seconds=120)
+
+        assert count == 1
+
+        await db_session.refresh(agent)
+        assert agent.state == "offline"
+        assert agent.status == "offline"
+
+    @pytest.mark.asyncio
+    async def test_keep_recent_agent_online(self, db_session, test_user):
+        """Should not mark agents with recent heartbeat as offline."""
+        from backend.services.agent_service import AgentService
+
+        # Create an agent with recent heartbeat
+        recent_time = datetime.now(UTC) - timedelta(seconds=30)
+        agent = Agent(
+            id=str(uuid.uuid4()),
+            user_id=test_user.id,
+            name="Active Agent",
+            state="online",
+            status="online",
+            last_heartbeat_at=recent_time,
+        )
+        db_session.add(agent)
+        await db_session.commit()
+
+        # Run staleness check with 2-minute threshold
+        svc = AgentService(db_session)
+        count = await svc.mark_stale_agents_offline(threshold_seconds=120)
+
+        assert count == 0
+
+        await db_session.refresh(agent)
+        assert agent.state == "online"
+
+    @pytest.mark.asyncio
+    async def test_mark_stale_agents_after_days(self, db_session, test_user):
+        """Should mark offline agents as stale after several days."""
+        from backend.services.agent_service import AgentService
+
+        # Create an agent offline for 10 days
+        old_time = datetime.now(UTC) - timedelta(days=10)
+        agent = Agent(
+            id=str(uuid.uuid4()),
+            user_id=test_user.id,
+            name="Old Offline Agent",
+            state="offline",
+            status="offline",
+            last_online_at=old_time,
+        )
+        db_session.add(agent)
+        await db_session.commit()
+
+        # Run staleness check with 7-day threshold
+        svc = AgentService(db_session)
+        count = await svc.mark_stale_agents(days=7)
+
+        assert count == 1
+
+        await db_session.refresh(agent)
+        assert agent.state == "stale"