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

AI seat tracking inserted as heartbeat into usage table.

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{}
 }

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

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)

coderd/database/dbmetrics/querymetrics.go

@@ -4792,6 +4792,14 @@ func (m queryMetricsStore) UpsertWorkspaceAppAuditSession(ctx context.Context, a
 	return r0, r1
 }
 
+func (m queryMetricsStore) UsageEventExistsByID(ctx context.Context, id string) (bool, error) {
+	start := time.Now()
+	r0, r1 := m.s.UsageEventExistsByID(ctx, id)
+	m.queryLatencies.WithLabelValues("UsageEventExistsByID").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UsageEventExistsByID").Inc()
+	return r0, r1
+}
+
 func (m queryMetricsStore) ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (database.ValidateGroupIDsRow, error) {
 	start := time.Now()
 	r0, r1 := m.s.ValidateGroupIDs(ctx, groupIds)

coderd/database/dbmock/dbmock.go

@@ -8945,6 +8945,21 @@ func (mr *MockStoreMockRecorder) UpsertWorkspaceAppAuditSession(ctx, arg any) *g
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertWorkspaceAppAuditSession", reflect.TypeOf((*MockStore)(nil).UpsertWorkspaceAppAuditSession), ctx, arg)
 }
 
+// UsageEventExistsByID mocks base method.
+func (m *MockStore) UsageEventExistsByID(ctx context.Context, id string) (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UsageEventExistsByID", ctx, id)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// UsageEventExistsByID indicates an expected call of UsageEventExistsByID.
+func (mr *MockStoreMockRecorder) UsageEventExistsByID(ctx, id any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UsageEventExistsByID", reflect.TypeOf((*MockStore)(nil).UsageEventExistsByID), ctx, id)
+}
+
 // ValidateGroupIDs mocks base method.
 func (m *MockStore) ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (database.ValidateGroupIDsRow, error) {
 	m.ctrl.T.Helper()

coderd/database/dump.sql

@@ -620,28 +620,35 @@ CREATE FUNCTION aggregate_usage_event() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
 BEGIN
-    -- Check for supported event types and throw error for unknown types
-    IF NEW.event_type NOT IN ('dc_managed_agents_v1') THEN
+    -- 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 (
-        -- Extract the date from the created_at timestamp, always using UTC for
-        -- consistency
         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
+            -- 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;
@@ -2655,7 +2662,7 @@ CREATE TABLE usage_events (
     publish_started_at timestamp with time zone,
     published_at timestamp with time zone,
     failure_message text,
-    CONSTRAINT usage_event_type_check CHECK ((event_type = 'dc_managed_agents_v1'::text))
+    CONSTRAINT usage_event_type_check CHECK ((event_type = ANY (ARRAY['dc_managed_agents_v1'::text, 'hb_ai_seats_v1'::text])))
 );
 
 COMMENT ON TABLE usage_events IS 'usage_events contains usage data that is collected from the product and potentially shipped to the usage collector service.';
@@ -3710,6 +3717,8 @@ CREATE INDEX idx_template_versions_has_ai_task ON template_versions USING btree
 
 CREATE UNIQUE INDEX idx_unique_preset_name ON template_version_presets USING btree (name, template_version_id);
 
+CREATE INDEX idx_usage_events_ai_seats ON usage_events USING btree (event_type, created_at) WHERE (event_type = 'hb_ai_seats_v1'::text);
+
 CREATE INDEX idx_usage_events_select_for_publishing ON usage_events USING btree (published_at, publish_started_at, created_at);
 
 CREATE INDEX idx_user_deleted_deleted_at ON user_deleted USING btree (deleted_at);

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
+);

coderd/database/querier.go

@@ -907,6 +907,7 @@ type sqlcQuerier interface {
 	// was started. This means that a new row was inserted (no previous session) or
 	// the updated_at is older than stale interval.
 	UpsertWorkspaceAppAuditSession(ctx context.Context, arg UpsertWorkspaceAppAuditSessionParams) (bool, error)
+	UsageEventExistsByID(ctx context.Context, id string) (bool, error)
 	ValidateGroupIDs(ctx context.Context, groupIds []uuid.UUID) (ValidateGroupIDsRow, error)
 	ValidateUserIDs(ctx context.Context, userIds []uuid.UUID) (ValidateUserIDsRow, error)
 }

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()
 

coderd/database/queries.sql.go

@@ -18845,6 +18845,19 @@ func (q *sqlQuerier) UpdateUsageEventsPostPublish(ctx context.Context, arg Updat
 	return err
 }
 
+const usageEventExistsByID = `-- name: UsageEventExistsByID :one
+SELECT EXISTS(
+    SELECT 1 FROM usage_events WHERE id = $1
+)::bool
+`
+
+func (q *sqlQuerier) UsageEventExistsByID(ctx context.Context, id string) (bool, error) {
+	row := q.db.QueryRowContext(ctx, usageEventExistsByID, id)
+	var column_1 bool
+	err := row.Scan(&column_1)
+	return column_1, err
+}
+
 const getUserLinkByLinkedID = `-- name: GetUserLinkByLinkedID :one
 SELECT
 	user_links.user_id, user_links.login_type, user_links.linked_id, user_links.oauth_access_token, user_links.oauth_refresh_token, user_links.oauth_expiry, user_links.oauth_access_token_key_id, user_links.oauth_refresh_token_key_id, user_links.claims

coderd/database/queries/usageevents.sql

@@ -15,6 +15,11 @@ VALUES
     (@id, @event_type, @event_data, @created_at, NULL, NULL, NULL)
 ON CONFLICT (id) DO NOTHING;
 
+-- name: UsageEventExistsByID :one
+SELECT EXISTS(
+    SELECT 1 FROM usage_events WHERE id = @id
+)::bool;
+
 -- name: SelectUsageEventsForPublishing :many
 WITH usage_events AS (
     UPDATE

coderd/pproflabel/pproflabel.go

@@ -34,6 +34,7 @@ const (
 	ServiceAgentMetricAggregator = "agent-metrics-aggregator"
 	// ServiceTallymanPublisher publishes usage events to coder/tallyman.
 	ServiceTallymanPublisher = "tallyman-publisher"
+	ServiceUsageEventCron    = "usage-event-cron"
 
 	RequestTypeTag = "coder_request_type"
 )

coderd/provisionerdserver/provisionerdserver_test.go

@@ -3020,7 +3020,7 @@ func TestCompleteJob(t *testing.T) {
 
 					// We never expect a usage event to be collected for
 					// template imports.
-					require.Empty(t, fakeUsageInserter.collectedEvents)
+					require.Equal(t, 0, fakeUsageInserter.TotalEventCount())
 				})
 			}
 		})
@@ -3371,13 +3371,13 @@ func TestCompleteJob(t *testing.T) {
 
 					if tc.expectUsageEvent {
 						// Check that a usage event was collected.
-						require.Len(t, fakeUsageInserter.collectedEvents, 1)
+						require.Len(t, fakeUsageInserter.GetDiscreteEvents(), 1)
 						require.Equal(t, usagetypes.DCManagedAgentsV1{
 							Count: 1,
-						}, fakeUsageInserter.collectedEvents[0])
+						}, fakeUsageInserter.GetDiscreteEvents()[0])
 					} else {
 						// Check that no usage event was collected.
-						require.Empty(t, fakeUsageInserter.collectedEvents)
+						require.Equal(t, 0, fakeUsageInserter.TotalEventCount())
 					}
 				})
 			}
@@ -5032,21 +5032,10 @@ func (s *fakeStream) cancel() {
 	s.c.Broadcast()
 }
 
-type fakeUsageInserter struct {
-	collectedEvents []usagetypes.Event
-}
-
-var _ usage.Inserter = &fakeUsageInserter{}
-
-func newFakeUsageInserter() (*fakeUsageInserter, *atomic.Pointer[usage.Inserter]) {
+func newFakeUsageInserter() (*coderdtest.UsageInserter, *atomic.Pointer[usage.Inserter]) {
 	poitr := &atomic.Pointer[usage.Inserter]{}
-	fake := &fakeUsageInserter{}
+	fake := coderdtest.NewUsageInserter()
 	var inserter usage.Inserter = fake
 	poitr.Store(&inserter)
 	return fake, poitr
 }
-
-func (f *fakeUsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
-	f.collectedEvents = append(f.collectedEvents, event)
-	return nil
-}

coderd/usage/inserter.go

@@ -14,6 +14,21 @@ type Inserter interface {
 	// The caller context must be authorized to create usage events in the
 	// database.
 	InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error
+
+	// InsertHeartbeatUsageEvent writes a heartbeat usage event to the database
+	// within the given transaction.
+	//
+	// The caller context must be authorized to create usage events in the database.
+	//
+	// The `id` should be a stable identifier for the event. Heartbeat events may be
+	// emitted by multiple replicas of the same daemon, so the same logical event
+	// may be submitted multiple times concurrently. For this reason the identifier
+	// must be deterministic and stateless, allowing duplicate submissions to be
+	// safely ignored.
+	//
+	// Inserts with the same `id` must be idempotent. The database enforces this by
+	// ignoring duplicate records.
+	InsertHeartbeatUsageEvent(ctx context.Context, tx database.Store, id string, event usagetypes.HeartbeatEvent) error
 }
 
 // AGPLInserter is a no-op implementation of Inserter.
@@ -30,3 +45,9 @@ func NewAGPLInserter() Inserter {
 func (AGPLInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, _ usagetypes.DiscreteEvent) error {
 	return nil
 }
+
+// InsertHeartbeatUsageEvent is a no-op implementation of
+// InsertHeartbeatUsageEvent.
+func (AGPLInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, _ string, _ usagetypes.HeartbeatEvent) error {
+	return nil
+}

coderd/usage/usagetypes/events.go

@@ -29,12 +29,15 @@ type UsageEventType string
 // ParseEventWithType function.
 const (
 	UsageEventTypeDCManagedAgentsV1 UsageEventType = "dc_managed_agents_v1"
+	UsageEventTypeHBAISeatsV1       UsageEventType = "hb_ai_seats_v1"
 )
 
 func (e UsageEventType) Valid() bool {
 	switch e {
 	case UsageEventTypeDCManagedAgentsV1:
 		return true
+	case UsageEventTypeHBAISeatsV1:
+		return true
 	default:
 		return false
 	}
@@ -96,6 +99,12 @@ func ParseEventWithType(eventType UsageEventType, data json.RawMessage) (Event,
 			return nil, err
 		}
 		return event, nil
+	case UsageEventTypeHBAISeatsV1:
+		var event HBAISeats
+		if err := ParseEvent(data, &event); err != nil {
+			return nil, err
+		}
+		return event, nil
 	default:
 		return nil, UnknownEventTypeError{EventType: string(eventType)}
 	}
@@ -121,6 +130,12 @@ type DiscreteEvent interface {
 	discreteUsageEvent() // marker method, also prevents external types from implementing this interface
 }
 
+// HeartbeatEvent is a usage event that is collected as a heartbeat.
+type HeartbeatEvent interface {
+	Event
+	heartbeatUsageEvent() // marker method, also prevents external types from implementing this interface
+}
+
 // DCManagedAgentsV1 is a discrete usage event for the number of managed agents.
 // This event is sent in the following situations:
 //   - Once on first startup after usage tracking is added to the product with
@@ -150,3 +165,30 @@ func (e DCManagedAgentsV1) Fields() map[string]any {
 		"count": e.Count,
 	}
 }
+
+// HBAISeats is a heartbeat event for the total number of AI seats consumed.
+type HBAISeats struct {
+	Count int64 `json:"count"`
+}
+
+var _ HeartbeatEvent = HBAISeats{}
+
+func (HBAISeats) usageEvent()          {}
+func (HBAISeats) heartbeatUsageEvent() {}
+func (HBAISeats) EventType() UsageEventType {
+	return UsageEventTypeHBAISeatsV1
+}
+
+func (e HBAISeats) Valid() error {
+	if e.Count < 0 {
+		return xerrors.New("count cannot be negative")
+	}
+	// The count can be 0
+	return nil
+}
+
+func (e HBAISeats) Fields() map[string]any {
+	return map[string]any{
+		"count": e.Count,
+	}
+}

coderd/usage/usagetypes/events_test.go

@@ -65,4 +65,15 @@ func TestParseEventWithType(t *testing.T) {
 		require.Equal(t, eventType, event.EventType())
 		require.Equal(t, map[string]any{"count": uint64(1)}, event.Fields())
 	})
+
+	t.Run("HBAISeatsV1", func(t *testing.T) {
+		t.Parallel()
+
+		eventType := usagetypes.UsageEventTypeHBAISeatsV1
+		event, err := usagetypes.ParseEventWithType(eventType, []byte(`{"count": 1}`))
+		require.NoError(t, err)
+		require.Equal(t, usagetypes.HBAISeats{Count: 1}, event)
+		require.Equal(t, eventType, event.EventType())
+		require.Equal(t, map[string]any{"count": int64(1)}, event.Fields())
+	})
 }