feat: track ai seat usage (#22682)

When a user uses an AI feature, we record them in the `ai_seat_state` as consuming a seat. 

Added in debouching to prevent excessive writes to the db for this feature. There is no need for frequent updates.

coderd/aiseats/aiseats.go

@@ -0,0 +1,38 @@
+// Package aiseats is the AGPL version the package.
+// The actual implementation is in `enterprise/aiseats`.
+package aiseats
+
+import (
+	"context"
+
+	"github.com/google/uuid"
+
+	"github.com/coder/coder/v2/coderd/database"
+)
+
+type Reason struct {
+	EventType   database.AiSeatUsageReason
+	Description string
+}
+
+// ReasonAIBridge constructs a reason for usage originating from AI Bridge.
+func ReasonAIBridge(description string) Reason {
+	return Reason{EventType: database.AiSeatUsageReasonAibridge, Description: description}
+}
+
+// ReasonTask constructs a reason for usage originating from tasks.
+func ReasonTask(description string) Reason {
+	return Reason{EventType: database.AiSeatUsageReasonTask, Description: description}
+}
+
+// SeatTracker records AI seat consumption state.
+type SeatTracker interface {
+	// RecordUsage does not return an error to prevent blocking the user from using
+	// AI features. This method is used to record usage, not enforce it.
+	RecordUsage(ctx context.Context, userID uuid.UUID, reason Reason)
+}
+
+// Noop is an AGPL seat tracker that does nothing.
+type Noop struct{}
+
+func (Noop) RecordUsage(context.Context, uuid.UUID, Reason) {}

coderd/coderd.go

@@ -44,6 +44,7 @@ import (
 	"github.com/coder/coder/v2/buildinfo"
 	"github.com/coder/coder/v2/coderd/agentapi"
 	"github.com/coder/coder/v2/coderd/agentapi/metadatabatcher"
+	"github.com/coder/coder/v2/coderd/aiseats"
 	_ "github.com/coder/coder/v2/coderd/apidoc" // Used for swagger docs.
 	"github.com/coder/coder/v2/coderd/appearance"
 	"github.com/coder/coder/v2/coderd/audit"
@@ -630,6 +631,8 @@ func New(options *Options) *API {
 		dbRolluper:       options.DatabaseRolluper,
 		ProfileCollector: defaultProfileCollector{},
 	}
+	api.AISeatTracker = aiseats.Noop{}
+
 	api.WorkspaceAppsProvider = workspaceapps.NewDBTokenProvider(
 		ctx,
 		options.Logger.Named("workspaceapps"),
@@ -2033,6 +2036,8 @@ type API struct {
 	dbRolluper *dbrollup.Rolluper
 	// chatDaemon handles background processing of pending chats.
 	chatDaemon *chatd.Server
+	// AISeatTracker records AI seat usage.
+	AISeatTracker aiseats.SeatTracker
 	// gitSyncWorker refreshes stale chat diff statuses in the
 	// background.
 	gitSyncWorker *gitsync.Worker
@@ -2245,6 +2250,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
 		provisionerdserver.Options{
 			OIDCConfig:          api.OIDCConfig,
 			ExternalAuthConfigs: api.ExternalAuthConfigs,
+			AISeatTracker:       api.AISeatTracker,
 			Clock:               api.Clock,
 			HeartbeatFn:         options.heartbeatFn,
 		},

coderd/provisionerdserver/provisionerdserver.go

@@ -28,6 +28,7 @@ import (
 	protobuf "google.golang.org/protobuf/proto"
 
 	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/coderd/aiseats"
 	"github.com/coder/coder/v2/coderd/apikey"
 	"github.com/coder/coder/v2/coderd/audit"
 	"github.com/coder/coder/v2/coderd/database"
@@ -76,6 +77,7 @@ const (
 type Options struct {
 	OIDCConfig          promoauth.OAuth2Config
 	ExternalAuthConfigs []*externalauth.Config
+	AISeatTracker       aiseats.SeatTracker
 
 	// Clock for testing
 	Clock quartz.Clock
@@ -120,6 +122,7 @@ type server struct {
 	NotificationsEnqueuer       notifications.Enqueuer
 	PrebuildsOrchestrator       *atomic.Pointer[prebuilds.ReconciliationOrchestrator]
 	UsageInserter               *atomic.Pointer[usage.Inserter]
+	AISeatTracker               aiseats.SeatTracker
 	Experiments                 codersdk.Experiments
 
 	OIDCConfig promoauth.OAuth2Config
@@ -215,6 +218,9 @@ func NewServer(
 	if err := tags.Valid(); err != nil {
 		return nil, xerrors.Errorf("invalid tags: %w", err)
 	}
+	if options.AISeatTracker == nil {
+		options.AISeatTracker = aiseats.Noop{}
+	}
 	if options.AcquireJobLongPollDur == 0 {
 		options.AcquireJobLongPollDur = DefaultAcquireJobLongPollDur
 	}
@@ -253,6 +259,7 @@ func NewServer(
 		heartbeatFn:                 options.HeartbeatFn,
 		PrebuildsOrchestrator:       prebuildsOrchestrator,
 		UsageInserter:               usageInserter,
+		AISeatTracker:               options.AISeatTracker,
 		metrics:                     metrics,
 		Experiments:                 experiments,
 	}
@@ -2437,6 +2444,12 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 		})
 	}
 
+	// Record AI seat usage for successful task workspace builds.
+	if workspaceBuild.Transition == database.WorkspaceTransitionStart && workspace.TaskID.Valid {
+		s.AISeatTracker.RecordUsage(ctx, workspace.OwnerID,
+			aiseats.ReasonTask("task workspace build succeeded"))
+	}
+
 	if s.PrebuildsOrchestrator != nil && input.PrebuiltWorkspaceBuildStage == sdkproto.PrebuiltWorkspaceBuildStage_CLAIM {
 		// Track resource replacements, if there are any.
 		orchestrator := s.PrebuildsOrchestrator.Load()