diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 2acc6178b8..95e5ce3a40 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2327,6 +2327,13 @@ func (q *querier) GetAPIKeysLastUsedAfter(ctx context.Context, lastUsed time.Tim return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetAPIKeysLastUsedAfter)(ctx, lastUsed) } +func (q *querier) GetActiveAISeatCount(ctx context.Context) (int64, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceLicense); err != nil { + return 0, err + } + return q.db.GetActiveAISeatCount(ctx) +} + func (q *querier) GetActivePresetPrebuildSchedules(ctx context.Context) ([]database.TemplateVersionPresetPrebuildSchedule, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTemplate.All()); err != nil { return nil, err @@ -6425,6 +6432,13 @@ func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg datab return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg) } +func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error { + if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceSystem); err != nil { + return err + } + return q.db.UpsertAISeatState(ctx, arg) +} + func (q *querier) UpsertAnnouncementBanners(ctx context.Context, value string) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index af412435bc..4a3bcf15f6 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1163,6 +1163,14 @@ func (s *MethodTestSuite) TestProvisionerJob() { } func (s *MethodTestSuite) TestLicense() { + s.Run("GetActiveAISeatCount", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + dbm.EXPECT().GetActiveAISeatCount(gomock.Any()).Return(int64(100), nil).AnyTimes() + check.Args().Asserts(rbac.ResourceLicense, policy.ActionRead).Returns(int64(100)) + })) + s.Run("UpsertAISeatState", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + dbm.EXPECT().UpsertAISeatState(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + check.Args(database.UpsertAISeatStateParams{}).Asserts(rbac.ResourceSystem, policy.ActionCreate) + })) s.Run("GetLicenses", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { a := database.License{ID: 1} b := database.License{ID: 2} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 59a1610a35..f6b021086b 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -871,6 +871,14 @@ func (m queryMetricsStore) GetAPIKeysLastUsedAfter(ctx context.Context, lastUsed return r0, r1 } +func (m queryMetricsStore) GetActiveAISeatCount(ctx context.Context) (int64, error) { + start := time.Now() + r0, r1 := m.s.GetActiveAISeatCount(ctx) + m.queryLatencies.WithLabelValues("GetActiveAISeatCount").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetActiveAISeatCount").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetActivePresetPrebuildSchedules(ctx context.Context) ([]database.TemplateVersionPresetPrebuildSchedule, error) { start := time.Now() r0, r1 := m.s.GetActivePresetPrebuildSchedules(ctx) @@ -4414,6 +4422,14 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context, return r0 } +func (m queryMetricsStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error { + start := time.Now() + r0 := m.s.UpsertAISeatState(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertAISeatState").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAISeatState").Inc() + return r0 +} + func (m queryMetricsStore) UpsertAnnouncementBanners(ctx context.Context, value string) error { start := time.Now() r0 := m.s.UpsertAnnouncementBanners(ctx, value) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index e0e4b3256e..bb56ae3e17 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1478,6 +1478,21 @@ func (mr *MockStoreMockRecorder) GetAPIKeysLastUsedAfter(ctx, lastUsed any) *gom return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAPIKeysLastUsedAfter", reflect.TypeOf((*MockStore)(nil).GetAPIKeysLastUsedAfter), ctx, lastUsed) } +// GetActiveAISeatCount mocks base method. +func (m *MockStore) GetActiveAISeatCount(ctx context.Context) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetActiveAISeatCount", ctx) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetActiveAISeatCount indicates an expected call of GetActiveAISeatCount. +func (mr *MockStoreMockRecorder) GetActiveAISeatCount(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetActiveAISeatCount", reflect.TypeOf((*MockStore)(nil).GetActiveAISeatCount), ctx) +} + // GetActivePresetPrebuildSchedules mocks base method. func (m *MockStore) GetActivePresetPrebuildSchedules(ctx context.Context) ([]database.TemplateVersionPresetPrebuildSchedule, error) { m.ctrl.T.Helper() @@ -8244,6 +8259,20 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspacesTTLByTemplateID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspacesTTLByTemplateID), ctx, arg) } +// UpsertAISeatState mocks base method. +func (m *MockStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertAISeatState", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpsertAISeatState indicates an expected call of UpsertAISeatState. +func (mr *MockStoreMockRecorder) UpsertAISeatState(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAISeatState", reflect.TypeOf((*MockStore)(nil).UpsertAISeatState), ctx, arg) +} + // UpsertAnnouncementBanners mocks base method. func (m *MockStore) UpsertAnnouncementBanners(ctx context.Context, value string) error { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 2da02f81c1..d1a7b91539 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -10,6 +10,11 @@ CREATE TYPE agent_key_scope_enum AS ENUM ( 'no_user_data' ); +CREATE TYPE ai_seat_usage_reason AS ENUM ( + 'aibridge', + 'task' +); + CREATE TYPE api_key_scope AS ENUM ( 'coder:all', 'coder:application_connect', @@ -1046,6 +1051,15 @@ BEGIN END; $$; +CREATE TABLE ai_seat_state ( + user_id uuid NOT NULL, + first_used_at timestamp with time zone NOT NULL, + last_used_at timestamp with time zone NOT NULL, + last_event_type ai_seat_usage_reason NOT NULL, + last_event_description text NOT NULL, + updated_at timestamp with time zone NOT NULL +); + CREATE TABLE aibridge_interceptions ( id uuid NOT NULL, initiator_id uuid NOT NULL, @@ -3156,6 +3170,9 @@ ALTER TABLE ONLY workspace_resource_metadata ALTER COLUMN id SET DEFAULT nextval ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); +ALTER TABLE ONLY ai_seat_state + ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); + ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); @@ -3816,6 +3833,9 @@ COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS 'U the uniqueness requirement. A trigger allows us to enforce uniqueness going forward without requiring a migration to clean up historical data.'; +ALTER TABLE ONLY ai_seat_state + ADD CONSTRAINT ai_seat_state_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; + ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id); diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index 2fb45a6963..cbb47ce680 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -6,6 +6,7 @@ type ForeignKeyConstraint string // ForeignKeyConstraint enums. const ( + ForeignKeyAiSeatStateUserID ForeignKeyConstraint = "ai_seat_state_user_id_fkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyAibridgeInterceptionsInitiatorID ForeignKeyConstraint = "aibridge_interceptions_initiator_id_fkey" // ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id); ForeignKeyAPIKeysUserIDUUID ForeignKeyConstraint = "api_keys_user_id_uuid_fkey" // ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyChatDiffStatusesChatID ForeignKeyConstraint = "chat_diff_statuses_chat_id_fkey" // ALTER TABLE ONLY chat_diff_statuses ADD CONSTRAINT chat_diff_statuses_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE; diff --git a/coderd/database/migrations/000439_ai_seat_state.down.sql b/coderd/database/migrations/000439_ai_seat_state.down.sql new file mode 100644 index 0000000000..aa9695366c --- /dev/null +++ b/coderd/database/migrations/000439_ai_seat_state.down.sql @@ -0,0 +1,3 @@ +DROP TABLE ai_seat_state; + +DROP TYPE ai_seat_usage_reason; diff --git a/coderd/database/migrations/000439_ai_seat_state.up.sql b/coderd/database/migrations/000439_ai_seat_state.up.sql new file mode 100644 index 0000000000..97efc68670 --- /dev/null +++ b/coderd/database/migrations/000439_ai_seat_state.up.sql @@ -0,0 +1,13 @@ +CREATE TYPE ai_seat_usage_reason AS ENUM ( + 'aibridge', + 'task' +); + +CREATE TABLE ai_seat_state ( + user_id uuid NOT NULL PRIMARY KEY REFERENCES users (id) ON DELETE CASCADE, + first_used_at timestamptz NOT NULL, + last_used_at timestamptz NOT NULL, + last_event_type ai_seat_usage_reason NOT NULL, + last_event_description text NOT NULL, + updated_at timestamptz NOT NULL +); diff --git a/coderd/database/migrations/testdata/fixtures/000439_ai_seat_state.up.sql b/coderd/database/migrations/testdata/fixtures/000439_ai_seat_state.up.sql new file mode 100644 index 0000000000..827697f7ee --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000439_ai_seat_state.up.sql @@ -0,0 +1,11 @@ +INSERT INTO + ai_seat_state ( + user_id, + first_used_at, + last_used_at, + last_event_type, + last_event_description, + updated_at + ) +VALUES + ('30095c71-380b-457a-8995-97b8ee6e5307', NOW(), NOW(), 'task'::ai_seat_usage_reason, 'Used for AI task', NOW()); diff --git a/coderd/database/models.go b/coderd/database/models.go index 93725ac787..9b07ae580a 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -741,6 +741,64 @@ func AllAgentKeyScopeEnumValues() []AgentKeyScopeEnum { } } +type AiSeatUsageReason string + +const ( + AiSeatUsageReasonAibridge AiSeatUsageReason = "aibridge" + AiSeatUsageReasonTask AiSeatUsageReason = "task" +) + +func (e *AiSeatUsageReason) Scan(src interface{}) error { + switch s := src.(type) { + case []byte: + *e = AiSeatUsageReason(s) + case string: + *e = AiSeatUsageReason(s) + default: + return fmt.Errorf("unsupported scan type for AiSeatUsageReason: %T", src) + } + return nil +} + +type NullAiSeatUsageReason struct { + AiSeatUsageReason AiSeatUsageReason `json:"ai_seat_usage_reason"` + Valid bool `json:"valid"` // Valid is true if AiSeatUsageReason is not NULL +} + +// Scan implements the Scanner interface. +func (ns *NullAiSeatUsageReason) Scan(value interface{}) error { + if value == nil { + ns.AiSeatUsageReason, ns.Valid = "", false + return nil + } + ns.Valid = true + return ns.AiSeatUsageReason.Scan(value) +} + +// Value implements the driver Valuer interface. +func (ns NullAiSeatUsageReason) Value() (driver.Value, error) { + if !ns.Valid { + return nil, nil + } + return string(ns.AiSeatUsageReason), nil +} + +func (e AiSeatUsageReason) Valid() bool { + switch e { + case AiSeatUsageReasonAibridge, + AiSeatUsageReasonTask: + return true + } + return false +} + +func AllAiSeatUsageReasonValues() []AiSeatUsageReason { + return []AiSeatUsageReason{ + AiSeatUsageReasonAibridge, + AiSeatUsageReasonTask, + } +} + type AppSharingLevel string const ( @@ -3975,6 +4033,15 @@ type APIKey struct { AllowList AllowList `db:"allow_list" json:"allow_list"` } +type AiSeatState struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + FirstUsedAt time.Time `db:"first_used_at" json:"first_used_at"` + LastUsedAt time.Time `db:"last_used_at" json:"last_used_at"` + LastEventType AiSeatUsageReason `db:"last_event_type" json:"last_event_type"` + LastEventDescription string `db:"last_event_description" json:"last_event_description"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` +} + type AuditLog struct { ID uuid.UUID `db:"id" json:"id"` Time time.Time `db:"time" json:"time"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 1aba92e04f..c252191483 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -187,6 +187,7 @@ type sqlcQuerier interface { GetAPIKeysByLoginType(ctx context.Context, arg GetAPIKeysByLoginTypeParams) ([]APIKey, error) GetAPIKeysByUserID(ctx context.Context, arg GetAPIKeysByUserIDParams) ([]APIKey, error) GetAPIKeysLastUsedAfter(ctx context.Context, lastUsed time.Time) ([]APIKey, error) + GetActiveAISeatCount(ctx context.Context) (int64, error) GetActivePresetPrebuildSchedules(ctx context.Context) ([]TemplateVersionPresetPrebuildSchedule, error) GetActiveUserCount(ctx context.Context, includeSystem bool) (int64, error) GetActiveWorkspaceBuildsByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceBuild, error) @@ -834,6 +835,7 @@ type sqlcQuerier interface { UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error + UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) error UpsertAnnouncementBanners(ctx context.Context, value string) error UpsertApplicationName(ctx context.Context, value string) error // Upserts boundary usage statistics for a replica. On INSERT (new period), uses diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3fcb18b891..79283ac791 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1198,6 +1198,64 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up return i, err } +const getActiveAISeatCount = `-- name: GetActiveAISeatCount :one +SELECT + COUNT(*) +FROM + ai_seat_state ais +JOIN + users u +ON + ais.user_id = u.id +WHERE + u.status = 'active'::user_status + AND u.deleted = false + AND u.is_system = false +` + +func (q *sqlQuerier) GetActiveAISeatCount(ctx context.Context) (int64, error) { + row := q.db.QueryRowContext(ctx, getActiveAISeatCount) + var count int64 + err := row.Scan(&count) + return count, err +} + +const upsertAISeatState = `-- name: UpsertAISeatState :exec +INSERT INTO ai_seat_state ( + user_id, + first_used_at, + last_used_at, + last_event_type, + last_event_description, + updated_at +) +VALUES + ($1, $2, $2, $3, $4, $2) +ON CONFLICT (user_id) DO UPDATE +SET + last_used_at = EXCLUDED.last_used_at, + last_event_type = EXCLUDED.last_event_type, + last_event_description = EXCLUDED.last_event_description, + updated_at = EXCLUDED.updated_at +` + +type UpsertAISeatStateParams struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + FirstUsedAt time.Time `db:"first_used_at" json:"first_used_at"` + LastEventType AiSeatUsageReason `db:"last_event_type" json:"last_event_type"` + LastEventDescription string `db:"last_event_description" json:"last_event_description"` +} + +func (q *sqlQuerier) UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) error { + _, err := q.db.ExecContext(ctx, upsertAISeatState, + arg.UserID, + arg.FirstUsedAt, + arg.LastEventType, + arg.LastEventDescription, + ) + return err +} + const deleteAPIKeyByID = `-- name: DeleteAPIKeyByID :exec DELETE FROM api_keys diff --git a/coderd/database/queries/aiseats.sql b/coderd/database/queries/aiseats.sql new file mode 100644 index 0000000000..628e2fffd6 --- /dev/null +++ b/coderd/database/queries/aiseats.sql @@ -0,0 +1,32 @@ +-- name: UpsertAISeatState :exec +INSERT INTO ai_seat_state ( + user_id, + first_used_at, + last_used_at, + last_event_type, + last_event_description, + updated_at +) +VALUES + ($1, $2, $2, $3, $4, $2) +ON CONFLICT (user_id) DO UPDATE +SET + last_used_at = EXCLUDED.last_used_at, + last_event_type = EXCLUDED.last_event_type, + last_event_description = EXCLUDED.last_event_description, + updated_at = EXCLUDED.updated_at +; + +-- name: GetActiveAISeatCount :one +SELECT + COUNT(*) +FROM + ai_seat_state ais +JOIN + users u +ON + ais.user_id = u.id +WHERE + u.status = 'active'::user_status + AND u.deleted = false + AND u.is_system = false; diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 380d2526b3..6066e4ea50 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -7,6 +7,7 @@ type UniqueConstraint string // UniqueConstraint enums. const ( UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); + UniqueAiSeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); UniqueAibridgeInterceptionsPkey UniqueConstraint = "aibridge_interceptions_pkey" // ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); UniqueAibridgeTokenUsagesPkey UniqueConstraint = "aibridge_token_usages_pkey" // ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); UniqueAibridgeToolUsagesPkey UniqueConstraint = "aibridge_tool_usages_pkey" // ALTER TABLE ONLY aibridge_tool_usages ADD CONSTRAINT aibridge_tool_usages_pkey PRIMARY KEY (id);