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

kairo_migrations/versions/001_add_agent_dashboard_extensions.py

@@ -0,0 +1,450 @@
+"""
+Add agent dashboard extensions
+
+Revision ID: 001_agent_dashboard
+Revises: (previous migration)
+Create Date: 2026-02-03
+
+This migration extends the existing agents table and creates new tables
+for the Agent Dashboard feature.
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers
+revision = '001_agent_dashboard'
+down_revision = None  # Set to previous migration ID
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # =========================================================================
+    # STEP 1: Create ENUM types
+    # =========================================================================
+
+    # Create enum types (if they don't exist)
+    agent_state_enum = postgresql.ENUM(
+        'created', 'online', 'idle', 'busy', 'offline', 'stale', 'disabled',
+        name='agent_state',
+        create_type=False,
+    )
+    agent_type_enum = postgresql.ENUM(
+        'sdk', 'container', 'simple',
+        name='agent_type',
+        create_type=False,
+    )
+    agent_event_type_enum = postgresql.ENUM(
+        'heartbeat', 'connected', 'disconnected', 'restart_requested',
+        'restart_completed', 'error', 'config_changed', 'state_changed',
+        'task_started', 'task_completed',
+        name='agent_event_type',
+        create_type=False,
+    )
+    alert_severity_enum = postgresql.ENUM(
+        'info', 'warning', 'error', 'critical',
+        name='alert_severity',
+        create_type=False,
+    )
+    alert_channel_enum = postgresql.ENUM(
+        'email', 'webhook', 'slack', 'pagerduty', 'in_app',
+        name='alert_channel',
+        create_type=False,
+    )
+    alert_status_enum = postgresql.ENUM(
+        'triggered', 'acknowledged', 'resolved', 'expired',
+        name='alert_status',
+        create_type=False,
+    )
+    setup_token_status_enum = postgresql.ENUM(
+        'pending', 'used', 'expired', 'revoked',
+        name='setup_token_status',
+        create_type=False,
+    )
+
+    # Create the enum types in the database
+    op.execute("CREATE TYPE agent_state AS ENUM ('created', 'online', 'idle', 'busy', 'offline', 'stale', 'disabled')")
+    op.execute("CREATE TYPE agent_type AS ENUM ('sdk', 'container', 'simple')")
+    op.execute("CREATE TYPE agent_event_type AS ENUM ('heartbeat', 'connected', 'disconnected', 'restart_requested', 'restart_completed', 'error', 'config_changed', 'state_changed', 'task_started', 'task_completed')")
+    op.execute("CREATE TYPE alert_severity AS ENUM ('info', 'warning', 'error', 'critical')")
+    op.execute("CREATE TYPE alert_channel AS ENUM ('email', 'webhook', 'slack', 'pagerduty', 'in_app')")
+    op.execute("CREATE TYPE alert_status AS ENUM ('triggered', 'acknowledged', 'resolved', 'expired')")
+    op.execute("CREATE TYPE setup_token_status AS ENUM ('pending', 'used', 'expired', 'revoked')")
+
+    # =========================================================================
+    # STEP 2: Extend existing agents table
+    # =========================================================================
+
+    # Add new columns to agents table
+    op.add_column('agents', sa.Column('state', agent_state_enum, nullable=True))
+    op.add_column('agents', sa.Column('previous_state', agent_state_enum, nullable=True))
+    op.add_column('agents', sa.Column('state_changed_at', sa.DateTime(timezone=True), nullable=True))
+    op.add_column('agents', sa.Column('last_heartbeat_at', sa.DateTime(timezone=True), 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('agent_type', agent_type_enum, nullable=True))
+    op.add_column('agents', sa.Column('sdk_version', sa.String(32), nullable=True))
+    op.add_column('agents', sa.Column('host_info', postgresql.JSONB, nullable=True))
+    op.add_column('agents', sa.Column('restart_count', sa.Integer, nullable=True, 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(255), nullable=True))
+    op.add_column('agents', sa.Column('config', postgresql.JSONB, nullable=True, server_default='{}'))
+    op.add_column('agents', sa.Column('capabilities', postgresql.JSONB, nullable=True, server_default='[]'))
+    op.add_column('agents', sa.Column('deleted_at', sa.DateTime(timezone=True), nullable=True))
+
+    # Set default values for existing rows
+    op.execute("UPDATE agents SET state = 'created' WHERE state IS NULL")
+    op.execute("UPDATE agents SET agent_type = 'sdk' WHERE agent_type IS NULL")
+    op.execute("UPDATE agents SET restart_count = 0 WHERE restart_count IS NULL")
+
+    # Make columns NOT NULL after setting defaults
+    op.alter_column('agents', 'state', nullable=False)
+    op.alter_column('agents', 'agent_type', nullable=False)
+    op.alter_column('agents', 'restart_count', nullable=False)
+
+    # Create indexes on agents table
+    op.create_index(
+        'idx_agents_state',
+        'agents',
+        ['state'],
+        postgresql_where=sa.text('deleted_at IS NULL'),
+    )
+    op.create_index(
+        'idx_agents_user_state',
+        'agents',
+        ['user_id', 'state'],
+        postgresql_where=sa.text('deleted_at IS NULL'),
+    )
+    op.create_index(
+        'idx_agents_last_heartbeat',
+        'agents',
+        [sa.text('last_heartbeat_at DESC')],
+        postgresql_where=sa.text('deleted_at IS NULL'),
+    )
+    op.create_index(
+        'idx_agents_agent_type',
+        'agents',
+        ['agent_type'],
+        postgresql_where=sa.text('deleted_at IS NULL'),
+    )
+    op.create_index(
+        'idx_agents_active',
+        'agents',
+        ['user_id', sa.text('last_heartbeat_at DESC')],
+        postgresql_where=sa.text("state IN ('online', 'idle', 'busy') AND deleted_at IS NULL"),
+    )
+
+    # =========================================================================
+    # STEP 3: Create agent_metrics_1m table
+    # =========================================================================
+
+    op.create_table(
+        'agent_metrics_1m',
+        sa.Column('id', sa.BigInteger, primary_key=True, autoincrement=True),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('bucket_time', sa.DateTime(timezone=True), nullable=False),
+        sa.Column('request_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('error_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('timeout_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('input_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('output_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('latency_sum_ms', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('latency_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('latency_min_ms', sa.Integer, nullable=True),
+        sa.Column('latency_max_ms', sa.Integer, nullable=True),
+        sa.Column('latency_p50_ms', sa.Integer, nullable=True),
+        sa.Column('latency_p99_ms', sa.Integer, nullable=True),
+        sa.Column('tool_call_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('tool_calls_by_name', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('model_usage', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    op.create_index('idx_metrics_1m_agent_time', 'agent_metrics_1m', ['agent_id', 'bucket_time'])
+    op.create_index('idx_metrics_1m_time', 'agent_metrics_1m', [sa.text('bucket_time DESC')])
+    op.create_unique_constraint('idx_metrics_1m_agent_bucket', 'agent_metrics_1m', ['agent_id', 'bucket_time'])
+
+    # =========================================================================
+    # STEP 4: Create agent_metrics_1h table
+    # =========================================================================
+
+    op.create_table(
+        'agent_metrics_1h',
+        sa.Column('id', sa.BigInteger, primary_key=True, autoincrement=True),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('bucket_time', sa.DateTime(timezone=True), nullable=False),
+        sa.Column('request_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('error_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('timeout_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('input_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('output_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('latency_sum_ms', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('latency_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('latency_min_ms', sa.Integer, nullable=True),
+        sa.Column('latency_max_ms', sa.Integer, nullable=True),
+        sa.Column('latency_p50_ms', sa.Integer, nullable=True),
+        sa.Column('latency_p99_ms', sa.Integer, nullable=True),
+        sa.Column('tool_call_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('tool_calls_by_name', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('model_usage', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('source_bucket_count', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    op.create_index('idx_metrics_1h_agent_time', 'agent_metrics_1h', ['agent_id', 'bucket_time'])
+    op.create_index('idx_metrics_1h_time', 'agent_metrics_1h', [sa.text('bucket_time DESC')])
+    op.create_unique_constraint('idx_metrics_1h_agent_bucket', 'agent_metrics_1h', ['agent_id', 'bucket_time'])
+
+    # =========================================================================
+    # STEP 5: Create agent_metrics_daily table
+    # =========================================================================
+
+    op.create_table(
+        'agent_metrics_daily',
+        sa.Column('id', sa.BigInteger, primary_key=True, autoincrement=True),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('bucket_date', sa.Date, nullable=False),
+        sa.Column('total_requests', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('total_errors', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('error_rate', sa.Numeric(5, 4), nullable=True),
+        sa.Column('total_input_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        sa.Column('total_output_tokens', sa.BigInteger, nullable=False, server_default='0'),
+        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('total_tool_calls', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('most_used_tools', postgresql.JSONB, nullable=False, server_default='[]'),
+        sa.Column('uptime_minutes', sa.Integer, nullable=False, server_default='0'),
+        sa.Column('uptime_percentage', sa.Numeric(5, 2), nullable=True),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    op.create_unique_constraint('idx_metrics_daily_agent_date', 'agent_metrics_daily', ['agent_id', 'bucket_date'])
+    op.create_index('idx_metrics_daily_date', 'agent_metrics_daily', [sa.text('bucket_date DESC')])
+
+    # =========================================================================
+    # STEP 6: Create agent_events table
+    # =========================================================================
+
+    op.create_table(
+        'agent_events',
+        sa.Column('id', sa.BigInteger, primary_key=True, autoincrement=True),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('event_type', agent_event_type_enum, nullable=False),
+        sa.Column('event_time', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('payload', postgresql.JSONB, nullable=True),
+        sa.Column('error_type', sa.String(64), nullable=True),
+        sa.Column('error_message', sa.Text, nullable=True),
+        sa.Column('severity', alert_severity_enum, nullable=True),
+        sa.Column('session_id', postgresql.UUID(as_uuid=True), nullable=True),
+    )
+
+    op.create_index('idx_events_agent_time', 'agent_events', ['agent_id', sa.text('event_time DESC')])
+    op.create_index('idx_events_type_time', 'agent_events', ['event_type', sa.text('event_time DESC')])
+    op.create_index('idx_events_agent_type', 'agent_events', ['agent_id', 'event_type', sa.text('event_time DESC')])
+    op.create_index(
+        'idx_events_session',
+        'agent_events',
+        ['session_id', 'event_time'],
+        postgresql_where=sa.text('session_id IS NOT NULL'),
+    )
+    op.create_index(
+        'idx_events_errors',
+        'agent_events',
+        ['agent_id', sa.text('event_time DESC')],
+        postgresql_where=sa.text("event_type = 'error'"),
+    )
+    op.create_index('idx_events_payload', 'agent_events', ['payload'], postgresql_using='gin')
+
+    # =========================================================================
+    # STEP 7: Create agent_alert_configs table
+    # =========================================================================
+
+    op.create_table(
+        'agent_alert_configs',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('uuid_generate_v4()')),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=True),
+        sa.Column('user_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('users.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('name', sa.String(128), nullable=False),
+        sa.Column('description', sa.Text, nullable=True),
+        sa.Column('enabled', sa.Boolean, nullable=False, server_default='true'),
+        sa.Column('metric_type', sa.String(64), nullable=False),
+        sa.Column('condition_operator', sa.String(8), nullable=False),
+        sa.Column('threshold_value', sa.Numeric, nullable=False),
+        sa.Column('threshold_unit', sa.String(16), nullable=True),
+        sa.Column('evaluation_window_minutes', sa.Integer, nullable=False, server_default='5'),
+        sa.Column('min_samples', sa.Integer, nullable=False, server_default='1'),
+        sa.Column('channels', postgresql.ARRAY(sa.String), nullable=False, server_default="'{in_app}'"),
+        sa.Column('notification_config', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('cooldown_minutes', sa.Integer, nullable=False, server_default='15'),
+        sa.Column('severity', alert_severity_enum, nullable=False, server_default="'warning'"),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.CheckConstraint("condition_operator IN ('gt', 'lt', 'gte', 'lte', 'eq')", name='valid_operator'),
+    )
+
+    op.create_index(
+        'idx_alert_configs_agent',
+        'agent_alert_configs',
+        ['agent_id'],
+        postgresql_where=sa.text('enabled = TRUE'),
+    )
+    op.create_index(
+        'idx_alert_configs_user',
+        'agent_alert_configs',
+        ['user_id'],
+        postgresql_where=sa.text('enabled = TRUE'),
+    )
+    op.create_index('idx_alert_configs_metric', 'agent_alert_configs', ['metric_type', 'enabled'])
+
+    # =========================================================================
+    # STEP 8: Create agent_alerts table
+    # =========================================================================
+
+    op.create_table(
+        'agent_alerts',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('uuid_generate_v4()')),
+        sa.Column('config_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agent_alert_configs.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('status', alert_status_enum, nullable=False, server_default="'triggered'"),
+        sa.Column('severity', alert_severity_enum, nullable=False),
+        sa.Column('triggered_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('triggered_value', sa.Numeric, nullable=False),
+        sa.Column('threshold_value', sa.Numeric, nullable=False),
+        sa.Column('title', sa.String(255), nullable=False),
+        sa.Column('message', sa.Text, nullable=True),
+        sa.Column('details', postgresql.JSONB, nullable=True),
+        sa.Column('acknowledged_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('acknowledged_by', postgresql.UUID(as_uuid=True), sa.ForeignKey('users.id'), nullable=True),
+        sa.Column('resolved_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('resolved_by', postgresql.UUID(as_uuid=True), sa.ForeignKey('users.id'), nullable=True),
+        sa.Column('resolution_notes', sa.Text, nullable=True),
+        sa.Column('auto_resolved', sa.Boolean, nullable=False, server_default='false'),
+        sa.Column('expires_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    op.create_index('idx_alerts_agent_status', 'agent_alerts', ['agent_id', 'status', sa.text('triggered_at DESC')])
+    op.create_index('idx_alerts_status_time', 'agent_alerts', ['status', sa.text('triggered_at DESC')])
+    op.create_index('idx_alerts_config', 'agent_alerts', ['config_id', sa.text('triggered_at DESC')])
+    op.create_index(
+        'idx_alerts_active',
+        'agent_alerts',
+        ['agent_id', sa.text('triggered_at DESC')],
+        postgresql_where=sa.text("status IN ('triggered', 'acknowledged')"),
+    )
+
+    # =========================================================================
+    # STEP 9: Create agent_setup_tokens table
+    # =========================================================================
+
+    op.create_table(
+        'agent_setup_tokens',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('uuid_generate_v4()')),
+        sa.Column('user_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('users.id', ondelete='CASCADE'), nullable=False),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=True),
+        sa.Column('token_hash', sa.String(128), nullable=False),
+        sa.Column('token_prefix', sa.String(8), nullable=False),
+        sa.Column('name', sa.String(128), nullable=True),
+        sa.Column('status', setup_token_status_enum, nullable=False, server_default="'pending'"),
+        sa.Column('expires_at', sa.DateTime(timezone=True), nullable=False),
+        sa.Column('used_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('used_from_ip', postgresql.INET, nullable=True),
+        sa.Column('used_agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id'), nullable=True),
+        sa.Column('initial_config', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+    )
+
+    op.create_index('idx_setup_tokens_prefix', 'agent_setup_tokens', ['token_prefix'])
+    op.create_index('idx_setup_tokens_user_status', 'agent_setup_tokens', ['user_id', 'status'])
+    op.create_index(
+        'idx_setup_tokens_agent',
+        'agent_setup_tokens',
+        ['agent_id'],
+        postgresql_where=sa.text('agent_id IS NOT NULL'),
+    )
+    op.create_index(
+        'idx_setup_tokens_active',
+        'agent_setup_tokens',
+        ['token_hash'],
+        postgresql_where=sa.text("status = 'pending'"),
+    )
+
+    # =========================================================================
+    # STEP 10: Create agent_pending_registrations table
+    # =========================================================================
+
+    op.create_table(
+        'agent_pending_registrations',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('uuid_generate_v4()')),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agents.id', ondelete='CASCADE'), nullable=False, unique=True),
+        sa.Column('setup_token_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('agent_setup_tokens.id', ondelete='SET NULL'), nullable=True),
+        sa.Column('expected_sdk_version', sa.String(32), nullable=True),
+        sa.Column('expected_capabilities', postgresql.ARRAY(sa.String), nullable=True),
+        sa.Column('wizard_step', sa.String(32), nullable=True),
+        sa.Column('wizard_data', postgresql.JSONB, nullable=False, server_default='{}'),
+        sa.Column('setup_instructions', sa.Text, nullable=True),
+        sa.Column('setup_command', sa.Text, nullable=True),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
+        sa.Column('expires_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
+    )
+
+    op.create_index('idx_pending_reg_agent', 'agent_pending_registrations', ['agent_id'])
+    op.create_index(
+        'idx_pending_reg_incomplete',
+        'agent_pending_registrations',
+        ['created_at'],
+        postgresql_where=sa.text('completed_at IS NULL'),
+    )
+
+
+def downgrade() -> None:
+    """Reverse the migration."""
+
+    # Drop tables in reverse order of creation
+    op.drop_table('agent_pending_registrations')
+    op.drop_table('agent_setup_tokens')
+    op.drop_table('agent_alerts')
+    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')
+
+    # Drop indexes on agents table
+    op.drop_index('idx_agents_active', table_name='agents')
+    op.drop_index('idx_agents_agent_type', table_name='agents')
+    op.drop_index('idx_agents_last_heartbeat', table_name='agents')
+    op.drop_index('idx_agents_user_state', table_name='agents')
+    op.drop_index('idx_agents_state', table_name='agents')
+
+    # Drop columns from agents table
+    op.drop_column('agents', 'deleted_at')
+    op.drop_column('agents', 'capabilities')
+    op.drop_column('agents', 'config')
+    op.drop_column('agents', 'last_restart_reason')
+    op.drop_column('agents', 'last_restart_at')
+    op.drop_column('agents', 'restart_count')
+    op.drop_column('agents', 'host_info')
+    op.drop_column('agents', 'sdk_version')
+    op.drop_column('agents', 'agent_type')
+    op.drop_column('agents', 'first_connected_at')
+    op.drop_column('agents', 'last_online_at')
+    op.drop_column('agents', 'last_heartbeat_at')
+    op.drop_column('agents', 'state_changed_at')
+    op.drop_column('agents', 'previous_state')
+    op.drop_column('agents', 'state')
+
+    # Drop enum types
+    op.execute('DROP TYPE setup_token_status')
+    op.execute('DROP TYPE alert_status')
+    op.execute('DROP TYPE alert_channel')
+    op.execute('DROP TYPE alert_severity')
+    op.execute('DROP TYPE agent_event_type')
+    op.execute('DROP TYPE agent_type')
+    op.execute('DROP TYPE agent_state')