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

AI seat tracking inserted as heartbeat into usage table.

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