coder · Emyrk · Mar 18, 2026 · Mar 5, 2026 · deansheather · Mar 13, 2026
diff --git a/coderd/coderdtest/usage.go b/coderd/coderdtest/usage.go
@@ -13,32 +13,64 @@ var _ usage.Inserter = (*UsageInserter)(nil)
 
 type UsageInserter struct {
 	sync.Mutex
-	events []usagetypes.DiscreteEvent
+	discreteEvents  []usagetypes.DiscreteEvent
+	heartbeatEvents []usagetypes.HeartbeatEvent
+	seenHeartbeats  map[string]struct{}
 }
 
 func NewUsageInserter() *UsageInserter {
 	return &UsageInserter{
-		events: []usagetypes.DiscreteEvent{},
+		discreteEvents:  []usagetypes.DiscreteEvent{},
+		seenHeartbeats:  map[string]struct{}{},
+		heartbeatEvents: []usagetypes.HeartbeatEvent{},
 	}
 }
 
 func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
 	u.Lock()
 	defer u.Unlock()
-	u.events = append(u.events, event)
+	u.discreteEvents = append(u.discreteEvents, event)
 	return nil
 }
 
-func (u *UsageInserter) GetEvents() []usagetypes.DiscreteEvent {
+func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error {
 	u.Lock()
 	defer u.Unlock()
-	eventsCopy := make([]usagetypes.DiscreteEvent, len(u.events))
-	copy(eventsCopy, u.events)
+	if _, seen := u.seenHeartbeats[id]; seen {
+		return nil
+	}
+
+	u.seenHeartbeats[id] = struct{}{}
+	u.heartbeatEvents = append(u.heartbeatEvents, event)
+	return nil
+}
+
+func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent {
+	u.Lock()
+	defer u.Unlock()
+	eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents))
+	copy(eventsCopy, u.heartbeatEvents)
+	return eventsCopy
+}
+
+func (u *UsageInserter) GetDiscreteEvents() []usagetypes.DiscreteEvent {
+	u.Lock()
+	defer u.Unlock()
+	eventsCopy := make([]usagetypes.DiscreteEvent, len(u.discreteEvents))
+	copy(eventsCopy, u.discreteEvents)
 	return eventsCopy
 }
 
+func (u *UsageInserter) TotalEventCount() int {
+	u.Lock()
+	defer u.Unlock()
+	return len(u.discreteEvents) + len(u.heartbeatEvents)
+}
+
 func (u *UsageInserter) Reset() {
 	u.Lock()
 	defer u.Unlock()
-	u.events = []usagetypes.DiscreteEvent{}
+	u.seenHeartbeats = map[string]struct{}{}
+	u.discreteEvents = []usagetypes.DiscreteEvent{}
+	u.heartbeatEvents = []usagetypes.HeartbeatEvent{}
 }
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -6805,6 +6805,13 @@ func (q *querier) UpsertWorkspaceAppAuditSession(ctx context.Context, arg databa
 	return q.db.UpsertWorkspaceAppAuditSession(ctx, arg)
 }
 
+func (q *querier) UsageEventExistsByID(ctx context.Context, id string) (bool, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceUsageEvent); err != nil {
+		return false, err
+	}
+	return q.db.UsageEventExistsByID(ctx, id)
+}
+
 func (q *querier) ValidateGroupIDs(ctx context.Context, groupIDs []uuid.UUID) (database.ValidateGroupIDsRow, error) {
 	// This check is probably overly restrictive, but the "correct" check isn't
 	// necessarily obvious. It's only used as a verification check for ACLs right

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -5123,6 +5123,12 @@ func (s *MethodTestSuite) TestUsageEvents() {
 		check.Args(params).Asserts(rbac.ResourceUsageEvent, policy.ActionCreate)
 	}))
 
+	s.Run("UsageEventExistsByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		id := uuid.NewString()
+		db.EXPECT().UsageEventExistsByID(gomock.Any(), id).Return(true, nil)
+		check.Args(id).Asserts(rbac.ResourceUsageEvent, policy.ActionRead)
+	}))
+
 	s.Run("SelectUsageEventsForPublishing", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		now := dbtime.Now()
 		db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), now).Return([]database.UsageEvent{}, nil)

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000444_usage_events_ai_seats.down.sql b/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;
diff --git a/coderd/database/migrations/000444_usage_events_ai_seats.up.sql b/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;
diff --git a/coderd/database/migrations/testdata/fixtures/000444_usage_events_ai_seats.up.sql b/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
+);
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -8315,6 +8315,80 @@ func TestUsageEventsTrigger(t *testing.T) {
 		require.WithinDuration(t, time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC), rows[1].Day, time.Second)
 	})
 
+	t.Run("HeartbeatAISeats", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
+
+		// Insert a heartbeat event.
+		err := db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        "hb-1",
+			EventType: "hb_ai_seats_v1",
+			EventData: []byte(`{"count": 10}`),
+			CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
+		})
+		require.NoError(t, err)
+
+		rows := getDailyRows(ctx, sqlDB)
+		require.Len(t, rows, 1)
+		require.Equal(t, "hb_ai_seats_v1", rows[0].EventType)
+		require.JSONEq(t, `{"count": 10}`, string(rows[0].UsageData))
+
+		// Insert a higher count on the same day — should take the max.
+		err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        "hb-2",
+			EventType: "hb_ai_seats_v1",
+			EventData: []byte(`{"count": 50}`),
+			CreatedAt: time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC),
+		})
+		require.NoError(t, err)
+
+		rows = getDailyRows(ctx, sqlDB)
+		require.Len(t, rows, 1)
+		require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
+
+		// Insert a lower count on the same day — should keep the max (50).
+		err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        "hb-3",
+			EventType: "hb_ai_seats_v1",
+			EventData: []byte(`{"count": 25}`),
+			CreatedAt: time.Date(2025, 1, 1, 18, 0, 0, 0, time.UTC),
+		})
+		require.NoError(t, err)
+
+		rows = getDailyRows(ctx, sqlDB)
+		require.Len(t, rows, 1)
+		require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
+
+		// Insert on a different day.
+		err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        "hb-4",
+			EventType: "hb_ai_seats_v1",
+			EventData: []byte(`{"count": 5}`),
+			CreatedAt: time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC),
+		})
+		require.NoError(t, err)
+
+		rows = getDailyRows(ctx, sqlDB)
+		require.Len(t, rows, 2)
+		require.JSONEq(t, `{"count": 50}`, string(rows[0].UsageData))
+		require.JSONEq(t, `{"count": 5}`, string(rows[1].UsageData))
+
+		// Also insert a dc_managed_agents_v1 on the same first day to
+		// verify different event types get separate daily rows.
+		err = db.InsertUsageEvent(ctx, database.InsertUsageEventParams{
+			ID:        "dc-1",
+			EventType: "dc_managed_agents_v1",
+			EventData: []byte(`{"count": 7}`),
+			CreatedAt: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
+		})
+		require.NoError(t, err)
+
+		rows = getDailyRows(ctx, sqlDB)
+		require.Len(t, rows, 3)
+	})
+
 	t.Run("UnknownEventType", func(t *testing.T) {
 		t.Parallel()