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

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

@@ -0,0 +1,349 @@
+"""
+End-to-end tests for Agent Setup Flow.
+
+Tests cover:
+- Generate setup token
+- Token expiry validation
+- Consume setup token successfully
+- Reject used/expired tokens
+- Agent transitions to "online" state after registration
+- API key is created with type="agent"
+"""
+
+import hashlib
+import pytest
+import uuid
+from datetime import datetime, timedelta, UTC
+from httpx import AsyncClient
+
+from backend.models.agent import Agent, AgentSetupToken
+from backend.models.api_key import ApiKey
+
+
+class TestGenerateSetupToken:
+    """Tests for POST /api/agents/{agent_id}/setup-token"""
+
+    @pytest.mark.asyncio
+    async def test_generate_setup_token_success(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should generate a valid setup token for an agent."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/setup-token", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "token" in data
+        assert data["token"].startswith("kairo_setup_")
+        assert "expires_at" in data
+        assert data["agent_id"] == test_agent.id
+
+    @pytest.mark.asyncio
+    async def test_setup_token_expiry_in_future(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should set token expiry in the future (default 15 minutes)."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/setup-token", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+
+        expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))
+        now = datetime.now(UTC)
+
+        # Token should expire within 10-20 minutes from now
+        assert expires_at > now + timedelta(minutes=10)
+        assert expires_at < now + timedelta(minutes=20)
+
+    @pytest.mark.asyncio
+    async def test_generate_token_for_nonexistent_agent(
+        self, client: AsyncClient, auth_headers: dict
+    ):
+        """Should return 404 for non-existent agent."""
+        fake_id = str(uuid.uuid4())
+
+        response = await client.post(
+            f"/api/agents/{fake_id}/setup-token", headers=auth_headers
+        )
+
+        assert response.status_code == 404
+
+    @pytest.mark.asyncio
+    async def test_generate_multiple_tokens(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should allow generating multiple tokens (old ones remain valid until used/expired)."""
+        response1 = await client.post(
+            f"/api/agents/{test_agent.id}/setup-token", headers=auth_headers
+        )
+        response2 = await client.post(
+            f"/api/agents/{test_agent.id}/setup-token", headers=auth_headers
+        )
+
+        assert response1.status_code == 200
+        assert response2.status_code == 200
+
+        token1 = response1.json()["token"]
+        token2 = response2.json()["token"]
+        assert token1 != token2
+
+
+class TestConsumeSetupToken:
+    """Tests for POST /api/agents/register"""
+
+    @pytest.mark.asyncio
+    async def test_consume_token_success(
+        self, client: AsyncClient, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should successfully register agent with valid token."""
+        # Generate a setup token
+        token, _ = await agent_factory.create_setup_token(test_agent.id)
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+            "host_info": {
+                "hostname": "test-machine",
+                "ip": "192.168.1.100",
+                "os": "linux",
+                "memory_mb": 8192,
+            },
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["agent_id"] == test_agent.id
+        assert data["name"] == test_agent.name
+        assert data["model_preference"] == test_agent.model_preference
+        assert "api_key" in data
+        assert data["api_key"].startswith("kairo_agent_")
+
+    @pytest.mark.asyncio
+    async def test_agent_state_transitions_to_online(
+        self, client: AsyncClient, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should transition agent state to 'online' after registration."""
+        token, _ = await agent_factory.create_setup_token(test_agent.id)
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+        }
+
+        await client.post("/api/agents/register", json=payload)
+
+        # Refresh agent from database
+        await db_session.refresh(test_agent)
+
+        assert test_agent.state == "online"
+        assert test_agent.first_connected_at is not None
+        assert test_agent.last_online_at is not None
+        assert test_agent.sdk_version == "1.0.0"
+
+    @pytest.mark.asyncio
+    async def test_api_key_created_with_agent_type(
+        self, client: AsyncClient, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should create API key with type='agent' after registration."""
+        from sqlalchemy import select
+
+        token, _ = await agent_factory.create_setup_token(test_agent.id)
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+        data = response.json()
+
+        # Find the created API key
+        stmt = select(ApiKey).where(ApiKey.agent_id == test_agent.id)
+        result = await db_session.execute(stmt)
+        api_key = result.scalar_one_or_none()
+
+        assert api_key is not None
+        assert api_key.key_type == "agent"
+        assert api_key.is_active is True
+        assert f"Agent: {test_agent.name}" in api_key.name
+
+    @pytest.mark.asyncio
+    async def test_reject_used_token(
+        self, client: AsyncClient, test_agent: Agent, agent_factory
+    ):
+        """Should reject a token that has already been used."""
+        # Create a used token
+        token, setup_token = await agent_factory.create_setup_token(
+            test_agent.id, used=True
+        )
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+
+        assert response.status_code == 401
+        assert "Invalid or expired" in response.json()["detail"]
+
+    @pytest.mark.asyncio
+    async def test_reject_expired_token(
+        self, client: AsyncClient, test_agent: Agent, db_session
+    ):
+        """Should reject a token that has expired."""
+        import secrets
+
+        # Create an expired token directly
+        token = f"kairo_setup_{secrets.token_urlsafe(32)}"
+        token_hash = hashlib.sha256(token.encode()).hexdigest()
+        expires_at = datetime.now(UTC) - timedelta(hours=1)  # Expired
+
+        setup_token = AgentSetupToken(
+            id=str(uuid.uuid4()),
+            agent_id=test_agent.id,
+            user_id=test_agent.user_id,
+            token_hash=token_hash,
+            expires_at=expires_at,
+            used=False,
+        )
+        db_session.add(setup_token)
+        await db_session.commit()
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+
+        assert response.status_code == 401
+        assert "Invalid or expired" in response.json()["detail"]
+
+    @pytest.mark.asyncio
+    async def test_reject_invalid_token(self, client: AsyncClient):
+        """Should reject an invalid token format."""
+        payload = {
+            "setup_token": "invalid_token_format",
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+
+        assert response.status_code == 401
+
+    @pytest.mark.asyncio
+    async def test_reject_nonexistent_token(self, client: AsyncClient):
+        """Should reject a token that does not exist in database."""
+        import secrets
+
+        payload = {
+            "setup_token": f"kairo_setup_{secrets.token_urlsafe(32)}",
+            "sdk_version": "1.0.0",
+        }
+
+        response = await client.post("/api/agents/register", json=payload)
+
+        assert response.status_code == 401
+
+    @pytest.mark.asyncio
+    async def test_token_marked_as_used_after_consumption(
+        self, client: AsyncClient, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should mark token as used after successful consumption."""
+        token, setup_token = await agent_factory.create_setup_token(test_agent.id)
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "1.0.0",
+        }
+
+        await client.post("/api/agents/register", json=payload)
+
+        # Refresh token from database
+        await db_session.refresh(setup_token)
+
+        assert setup_token.used is True
+        assert setup_token.used_at is not None
+
+    @pytest.mark.asyncio
+    async def test_host_info_stored_on_agent(
+        self, client: AsyncClient, test_agent: Agent, agent_factory, db_session
+    ):
+        """Should store host_info on the agent after registration."""
+        token, _ = await agent_factory.create_setup_token(test_agent.id)
+
+        host_info = {
+            "hostname": "production-server",
+            "ip": "10.0.0.50",
+            "os": "ubuntu-22.04",
+            "memory_mb": 16384,
+        }
+
+        payload = {
+            "setup_token": token,
+            "sdk_version": "2.0.0",
+            "host_info": host_info,
+        }
+
+        await client.post("/api/agents/register", json=payload)
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.host_info is not None
+        assert test_agent.host_info["hostname"] == "production-server"
+        assert test_agent.host_info["os"] == "ubuntu-22.04"
+
+
+class TestConnectionStatus:
+    """Tests for GET /api/agents/{agent_id}/connection-status"""
+
+    @pytest.mark.asyncio
+    async def test_connection_status_not_connected(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should return connected=false for unregistered agent."""
+        response = await client.get(
+            f"/api/agents/{test_agent.id}/connection-status", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["connected"] is False
+        assert data["state"] == "created"
+        assert data["first_connected_at"] is None
+
+    @pytest.mark.asyncio
+    async def test_connection_status_connected(
+        self, client: AsyncClient, auth_headers: dict, online_agent: Agent
+    ):
+        """Should return connected=true for registered online agent."""
+        response = await client.get(
+            f"/api/agents/{online_agent.id}/connection-status", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["connected"] is True
+        assert data["state"] == "online"
+        assert data["first_connected_at"] is not None
+        assert data["sdk_version"] == "1.0.0"
+        assert data["host_info"] is not None
+
+    @pytest.mark.asyncio
+    async def test_connection_status_not_found(
+        self, client: AsyncClient, auth_headers: dict
+    ):
+        """Should return 404 for non-existent agent."""
+        fake_id = str(uuid.uuid4())
+
+        response = await client.get(
+            f"/api/agents/{fake_id}/connection-status", headers=auth_headers
+        )
+
+        assert response.status_code == 404

kairo/migrations/versions/010_agent_dashboard.py

@@ -0,0 +1,246 @@
+"""Add agent dashboard tables and extend agents
+
+Revision ID: 010
+Revises: 009
+Create Date: 2026-02-03 00:00:00.000000
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import JSONB
+
+revision = "010"
+down_revision = "009"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Extend agents table with dashboard fields
+    op.add_column("agents", sa.Column("state", sa.String(), server_default="created"))
+    op.add_column("agents", sa.Column("agent_type", sa.String(), server_default="sdk"))
+    op.add_column("agents", sa.Column("sdk_version", sa.String(), nullable=True))
+    op.add_column("agents", sa.Column("host_info", JSONB, nullable=True))
+    op.add_column("agents", sa.Column("last_online_at", sa.DateTime(timezone=True), nullable=True))
+    op.add_column("agents", sa.Column("first_connected_at", sa.DateTime(timezone=True), nullable=True))
+    op.add_column("agents", sa.Column("restart_count", sa.Integer(), server_default="0"))
+    op.add_column("agents", sa.Column("last_restart_at", sa.DateTime(timezone=True), nullable=True))
+    op.add_column("agents", sa.Column("last_restart_reason", sa.String(), nullable=True))
+    op.add_column("agents", sa.Column("last_error_at", sa.DateTime(timezone=True), nullable=True))
+    op.add_column("agents", sa.Column("last_error_message", sa.Text(), nullable=True))
+    op.add_column("agents", sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True))
+
+    op.create_index("ix_agents_state", "agents", ["state"])
+    op.create_index("ix_agents_user_state", "agents", ["user_id", "state"])
+
+    # Add agent_id to api_keys for agent-scoped keys
+    op.add_column("api_keys", sa.Column(
+        "agent_id",
+        sa.String(),
+        sa.ForeignKey("agents.id", ondelete="CASCADE"),
+        nullable=True
+    ))
+    op.create_index("ix_api_keys_agent_id", "api_keys", ["agent_id"])
+
+    # Agent metrics - 1 minute granularity (30 day retention)
+    op.create_table(
+        "agent_metrics_1m",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("bucket_time", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("request_count", sa.Integer(), server_default="0"),
+        sa.Column("error_count", sa.Integer(), server_default="0"),
+        sa.Column("timeout_count", sa.Integer(), server_default="0"),
+        sa.Column("input_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("output_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("total_latency_ms", sa.BigInteger(), server_default="0"),
+        sa.Column("min_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("max_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("tool_calls", sa.Integer(), server_default="0"),
+        sa.Column("tool_errors", sa.Integer(), server_default="0"),
+        sa.Column("tool_breakdown", JSONB, nullable=True),
+        sa.Column("model_breakdown", JSONB, nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_metrics_1m_agent_time", "agent_metrics_1m", ["agent_id", "bucket_time"])
+    op.create_index("ix_agent_metrics_1m_bucket_time", "agent_metrics_1m", ["bucket_time"])
+
+    # Agent metrics - 1 hour granularity (1 year retention)
+    op.create_table(
+        "agent_metrics_1h",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("bucket_time", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("request_count", sa.Integer(), server_default="0"),
+        sa.Column("error_count", sa.Integer(), server_default="0"),
+        sa.Column("timeout_count", sa.Integer(), server_default="0"),
+        sa.Column("input_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("output_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("total_latency_ms", sa.BigInteger(), server_default="0"),
+        sa.Column("min_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("max_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("p50_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("p99_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("tool_calls", sa.Integer(), server_default="0"),
+        sa.Column("tool_errors", sa.Integer(), server_default="0"),
+        sa.Column("tool_breakdown", JSONB, nullable=True),
+        sa.Column("model_breakdown", JSONB, nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_metrics_1h_agent_time", "agent_metrics_1h", ["agent_id", "bucket_time"])
+    op.create_index("ix_agent_metrics_1h_bucket_time", "agent_metrics_1h", ["bucket_time"])
+
+    # Agent metrics - daily granularity (2 year retention)
+    op.create_table(
+        "agent_metrics_daily",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("date", sa.Date(), nullable=False),
+        sa.Column("request_count", sa.Integer(), server_default="0"),
+        sa.Column("error_count", sa.Integer(), server_default="0"),
+        sa.Column("timeout_count", sa.Integer(), server_default="0"),
+        sa.Column("input_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("output_tokens", sa.BigInteger(), server_default="0"),
+        sa.Column("total_latency_ms", sa.BigInteger(), server_default="0"),
+        sa.Column("min_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("max_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("avg_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("p50_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("p99_latency_ms", sa.Integer(), nullable=True),
+        sa.Column("tool_calls", sa.Integer(), server_default="0"),
+        sa.Column("tool_errors", sa.Integer(), server_default="0"),
+        sa.Column("tool_breakdown", JSONB, nullable=True),
+        sa.Column("model_breakdown", JSONB, nullable=True),
+        sa.Column("uptime_seconds", sa.Integer(), server_default="0"),
+        sa.Column("downtime_seconds", sa.Integer(), server_default="0"),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_metrics_daily_agent_date", "agent_metrics_daily", ["agent_id", "date"], unique=True)
+    op.create_index("ix_agent_metrics_daily_date", "agent_metrics_daily", ["date"])
+
+    # Agent events log
+    op.create_table(
+        "agent_events",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("event_type", sa.String(), nullable=False),
+        sa.Column("event_data", JSONB, nullable=True),
+        sa.Column("error_type", sa.String(), nullable=True),
+        sa.Column("error_message", sa.Text(), nullable=True),
+        sa.Column("session_id", sa.String(), nullable=True),
+        sa.Column("client_ip", sa.String(), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_events_agent_time", "agent_events", ["agent_id", "created_at"])
+    op.create_index("ix_agent_events_agent_type", "agent_events", ["agent_id", "event_type"])
+    op.create_index("ix_agent_events_created_at", "agent_events", ["created_at"])
+
+    # Agent alert configurations
+    op.create_table(
+        "agent_alert_configs",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("user_id", sa.String(), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("name", sa.String(), nullable=False),
+        sa.Column("alert_type", sa.String(), nullable=False),
+        sa.Column("metric", sa.String(), nullable=False),
+        sa.Column("condition", sa.String(), nullable=False),
+        sa.Column("threshold", sa.Float(), nullable=False),
+        sa.Column("window_seconds", sa.Integer(), server_default="300"),
+        sa.Column("cooldown_seconds", sa.Integer(), server_default="3600"),
+        sa.Column("severity", sa.String(), server_default="'warning'"),
+        sa.Column("channels", JSONB, server_default="'[\"email\"]'"),
+        sa.Column("is_enabled", sa.Boolean(), server_default="true"),
+        sa.Column("last_triggered_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_alert_configs_agent", "agent_alert_configs", ["agent_id"])
+    op.create_index("ix_agent_alert_configs_user", "agent_alert_configs", ["user_id"])
+
+    # Agent alert history
+    op.create_table(
+        "agent_alert_history",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("config_id", sa.String(), sa.ForeignKey("agent_alert_configs.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("alert_type", sa.String(), nullable=False),
+        sa.Column("severity", sa.String(), nullable=False),
+        sa.Column("status", sa.String(), server_default="'triggered'"),
+        sa.Column("message", sa.Text(), nullable=True),
+        sa.Column("trigger_value", sa.Float(), nullable=True),
+        sa.Column("threshold_value", sa.Float(), nullable=True),
+        sa.Column("acknowledged_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("acknowledged_by", sa.String(), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
+        sa.Column("resolved_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_alert_history_agent", "agent_alert_history", ["agent_id"])
+    op.create_index("ix_agent_alert_history_config", "agent_alert_history", ["config_id"])
+    op.create_index("ix_agent_alert_history_status", "agent_alert_history", ["status"])
+    op.create_index("ix_agent_alert_history_created", "agent_alert_history", ["created_at"])
+
+    # Agent setup tokens (one-time tokens for guided wizard)
+    op.create_table(
+        "agent_setup_tokens",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("user_id", sa.String(), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("token_hash", sa.String(), nullable=False, unique=True),
+        sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("used", sa.Boolean(), server_default="false"),
+        sa.Column("used_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("used_from_ip", sa.String(), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_setup_tokens_agent", "agent_setup_tokens", ["agent_id"])
+    op.create_index("ix_agent_setup_tokens_token_hash", "agent_setup_tokens", ["token_hash"], unique=True)
+
+    # Agent pending commands (for restart/stop via heartbeat)
+    op.create_table(
+        "agent_commands",
+        sa.Column("id", sa.String(), primary_key=True),
+        sa.Column("agent_id", sa.String(), sa.ForeignKey("agents.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("command_type", sa.String(), nullable=False),
+        sa.Column("payload", JSONB, nullable=True),
+        sa.Column("issued_by", sa.String(), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
+        sa.Column("status", sa.String(), server_default="'pending'"),
+        sa.Column("dispatched_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("acknowledged_at", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
+    )
+    op.create_index("ix_agent_commands_agent_status", "agent_commands", ["agent_id", "status"])
+    op.create_index("ix_agent_commands_expires", "agent_commands", ["expires_at"])
+
+
+def downgrade() -> None:
+    # Drop all new tables
+    op.drop_table("agent_commands")
+    op.drop_table("agent_setup_tokens")
+    op.drop_table("agent_alert_history")
+    op.drop_table("agent_alert_configs")
+    op.drop_table("agent_events")
+    op.drop_table("agent_metrics_daily")
+    op.drop_table("agent_metrics_1h")
+    op.drop_table("agent_metrics_1m")
+
+    # Remove api_keys.agent_id
+    op.drop_index("ix_api_keys_agent_id", table_name="api_keys")
+    op.drop_column("api_keys", "agent_id")
+
+    # Remove agents columns
+    op.drop_index("ix_agents_user_state", table_name="agents")
+    op.drop_index("ix_agents_state", table_name="agents")
+    op.drop_column("agents", "deleted_at")
+    op.drop_column("agents", "last_error_message")
+    op.drop_column("agents", "last_error_at")
+    op.drop_column("agents", "last_restart_reason")
+    op.drop_column("agents", "last_restart_at")
+    op.drop_column("agents", "restart_count")
+    op.drop_column("agents", "first_connected_at")
+    op.drop_column("agents", "last_online_at")
+    op.drop_column("agents", "host_info")
+    op.drop_column("agents", "sdk_version")
+    op.drop_column("agents", "agent_type")
+    op.drop_column("agents", "state")

{kairo_code-0.1.0.dist-info → kairo_code-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kairo-code
-Version: 0.1.0
+Version: 0.2.0
 Summary: Kairo Code - AI Coding Assistant by Kairon Labs
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown