chore: add tallyman events for ai seat tracking (#22689)

AI seat tracking inserted as heartbeat into usage table.

coderd/database/migrations/000444_usage_events_ai_seats.down.sql

@@ -0,0 +1,38 @@
+DROP INDEX IF EXISTS idx_usage_events_ai_seats;
+
+-- Remove hb_ai_seats_v1 rows so the original constraint can be restored.
+DELETE FROM usage_events WHERE event_type = 'hb_ai_seats_v1';
+DELETE FROM usage_events_daily WHERE event_type = 'hb_ai_seats_v1';
+
+-- Restore original constraint.
+ALTER TABLE usage_events
+  DROP CONSTRAINT usage_event_type_check,
+  ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1'));
+
+-- Restore the original aggregate function without hb_ai_seats_v1 support.
+CREATE OR REPLACE FUNCTION aggregate_usage_event()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF NEW.event_type NOT IN ('dc_managed_agents_v1') THEN
+        RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
+    END IF;
+
+    INSERT INTO usage_events_daily (day, event_type, usage_data)
+    VALUES (
+        date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date,
+        NEW.event_type,
+        NEW.event_data
+    )
+    ON CONFLICT (day, event_type) DO UPDATE SET
+        usage_data = CASE
+            WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN
+                jsonb_build_object(
+                    'count',
+                    COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) +
+                    COALESCE((NEW.event_data->>'count')::bigint, 0)
+                )
+        END;
+
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;

coderd/database/migrations/000444_usage_events_ai_seats.up.sql

@@ -0,0 +1,50 @@
+-- Expand the CHECK constraint to allow hb_ai_seats_v1.
+ALTER TABLE usage_events
+  DROP CONSTRAINT usage_event_type_check,
+  ADD CONSTRAINT usage_event_type_check CHECK (event_type IN ('dc_managed_agents_v1', 'hb_ai_seats_v1'));
+
+-- Partial index for efficient lookups of AI seat heartbeat events by time.
+-- This will be used for the admin dashboard to see seat count over time.
+CREATE INDEX idx_usage_events_ai_seats
+  ON usage_events (event_type, created_at)
+  WHERE event_type = 'hb_ai_seats_v1';
+
+-- Update the aggregate function to handle hb_ai_seats_v1 events.
+-- Heartbeat events replace the previous value for the same time period.
+CREATE OR REPLACE FUNCTION aggregate_usage_event()
+RETURNS TRIGGER AS $$
+BEGIN
+    -- Check for supported event types and throw error for unknown types.
+    IF NEW.event_type NOT IN ('dc_managed_agents_v1', 'hb_ai_seats_v1') THEN
+        RAISE EXCEPTION 'Unhandled usage event type in aggregate_usage_event: %', NEW.event_type;
+    END IF;
+
+    INSERT INTO usage_events_daily (day, event_type, usage_data)
+    VALUES (
+        date_trunc('day', NEW.created_at AT TIME ZONE 'UTC')::date,
+        NEW.event_type,
+        NEW.event_data
+    )
+    ON CONFLICT (day, event_type) DO UPDATE SET
+        usage_data = CASE
+            -- Handle simple counter events by summing the count.
+            WHEN NEW.event_type IN ('dc_managed_agents_v1') THEN
+                jsonb_build_object(
+                    'count',
+                    COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0) +
+                    COALESCE((NEW.event_data->>'count')::bigint, 0)
+                )
+			-- Heartbeat events: keep the max value seen that day
+            WHEN NEW.event_type IN ('hb_ai_seats_v1') THEN
+				jsonb_build_object(
+					'count',
+					GREATEST(
+						COALESCE((usage_events_daily.usage_data->>'count')::bigint, 0),
+						COALESCE((NEW.event_data->>'count')::bigint, 0)
+					)
+				)
+        END;
+
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;

coderd/database/migrations/testdata/fixtures/000444_usage_events_ai_seats.up.sql

@@ -0,0 +1,20 @@
+INSERT INTO usage_events (
+    id,
+    event_type,
+    event_data,
+    created_at,
+    publish_started_at,
+    published_at,
+    failure_message
+)
+VALUES
+-- Unpublished hb_ai_seats_v1 event.
+(
+    'ai-seats-event1',
+    'hb_ai_seats_v1',
+    '{"count":3}',
+    '2023-06-01 00:00:00+00',
+    NULL,
+    NULL,
+    NULL
+);