diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index ab600b6737..a4afb616e6 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -18478,7 +18478,8 @@ const docTemplate = `{ "idp_sync_settings_role", "workspace_agent", "workspace_app", - "task" + "task", + "ai_seat" ], "x-enum-varnames": [ "ResourceTypeTemplate", @@ -18506,7 +18507,8 @@ const docTemplate = `{ "ResourceTypeIdpSyncSettingsRole", "ResourceTypeWorkspaceAgent", "ResourceTypeWorkspaceApp", - "ResourceTypeTask" + "ResourceTypeTask", + "ResourceTypeAISeat" ] }, "codersdk.Response": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index ea32d716ad..99dbcd6c51 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -16871,7 +16871,8 @@ "idp_sync_settings_role", "workspace_agent", "workspace_app", - "task" + "task", + "ai_seat" ], "x-enum-varnames": [ "ResourceTypeTemplate", @@ -16899,7 +16900,8 @@ "ResourceTypeIdpSyncSettingsRole", "ResourceTypeWorkspaceAgent", "ResourceTypeWorkspaceApp", - "ResourceTypeTask" + "ResourceTypeTask", + "ResourceTypeAISeat" ] }, "codersdk.Response": { diff --git a/coderd/audit/diff.go b/coderd/audit/diff.go index c14dbc392f..e085c7d9ea 100644 --- a/coderd/audit/diff.go +++ b/coderd/audit/diff.go @@ -32,7 +32,8 @@ type Auditable interface { idpsync.OrganizationSyncSettings | idpsync.GroupSyncSettings | idpsync.RoleSyncSettings | - database.TaskTable + database.TaskTable | + database.AiSeatState } // Map is a map of changed fields in an audited resource. It maps field names to diff --git a/coderd/audit/request.go b/coderd/audit/request.go index 0c9a4bc4a2..147e53e4f7 100644 --- a/coderd/audit/request.go +++ b/coderd/audit/request.go @@ -132,6 +132,8 @@ func ResourceTarget[T Auditable](tgt T) string { return "Organization Role Sync" case database.TaskTable: return typed.Name + case database.AiSeatState: + return "AI Seat" default: panic(fmt.Sprintf("unknown resource %T for ResourceTarget", tgt)) } @@ -196,6 +198,8 @@ func ResourceID[T Auditable](tgt T) uuid.UUID { return noID // Org field on audit log has org id case database.TaskTable: return typed.ID + case database.AiSeatState: + return typed.UserID default: panic(fmt.Sprintf("unknown resource %T for ResourceID", tgt)) } @@ -251,6 +255,8 @@ func ResourceType[T Auditable](tgt T) database.ResourceType { return database.ResourceTypeIdpSyncSettingsGroup case database.TaskTable: return database.ResourceTypeTask + case database.AiSeatState: + return database.ResourceTypeAiSeat default: panic(fmt.Sprintf("unknown resource %T for ResourceType", typed)) } @@ -309,6 +315,8 @@ func ResourceRequiresOrgID[T Auditable]() bool { return true case database.TaskTable: return true + case database.AiSeatState: + return false default: panic(fmt.Sprintf("unknown resource %T for ResourceRequiresOrgID", tgt)) } diff --git a/coderd/coderd.go b/coderd/coderd.go index 4a3f5cd1ab..f6bdf09d2a 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -630,8 +630,8 @@ func New(options *Options) *API { ), dbRolluper: options.DatabaseRolluper, ProfileCollector: defaultProfileCollector{}, + AISeatTracker: aiseats.Noop{}, } - api.AISeatTracker = aiseats.Noop{} api.WorkspaceAppsProvider = workspaceapps.NewDBTokenProvider( ctx, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 95e5ce3a40..43f202d063 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -6432,9 +6432,9 @@ 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 { +func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceSystem); err != nil { - return err + return false, err } return q.db.UpsertAISeatState(ctx, arg) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 4a3bcf15f6..c5b5070ef7 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1168,7 +1168,7 @@ func (s *MethodTestSuite) TestLicense() { 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() + dbm.EXPECT().UpsertAISeatState(gomock.Any(), gomock.Any()).Return(true, nil).AnyTimes() check.Args(database.UpsertAISeatStateParams{}).Asserts(rbac.ResourceSystem, policy.ActionCreate) })) s.Run("GetLicenses", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index f6b021086b..f1b275f23c 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -4422,12 +4422,12 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context, return r0 } -func (m queryMetricsStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error { +func (m queryMetricsStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { start := time.Now() - r0 := m.s.UpsertAISeatState(ctx, arg) + r0, r1 := 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 + return r0, r1 } func (m queryMetricsStore) UpsertAnnouncementBanners(ctx context.Context, value string) error { diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index bb56ae3e17..1c33c22eeb 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -8260,11 +8260,12 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) * } // UpsertAISeatState mocks base method. -func (m *MockStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error { +func (m *MockStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "UpsertAISeatState", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 } // UpsertAISeatState indicates an expected call of UpsertAISeatState. diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index d1a7b91539..e81cf62f72 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -508,7 +508,8 @@ CREATE TYPE resource_type AS ENUM ( 'workspace_agent', 'workspace_app', 'prebuilds_settings', - 'task' + 'task', + 'ai_seat' ); CREATE TYPE startup_script_behavior AS ENUM ( diff --git a/coderd/database/migrations/000440_ai_seat_audit.down.sql b/coderd/database/migrations/000440_ai_seat_audit.down.sql new file mode 100644 index 0000000000..549da373b6 --- /dev/null +++ b/coderd/database/migrations/000440_ai_seat_audit.down.sql @@ -0,0 +1 @@ +-- resource_type enum values cannot be removed safely; no-op. diff --git a/coderd/database/migrations/000440_ai_seat_audit.up.sql b/coderd/database/migrations/000440_ai_seat_audit.up.sql new file mode 100644 index 0000000000..1728b30104 --- /dev/null +++ b/coderd/database/migrations/000440_ai_seat_audit.up.sql @@ -0,0 +1 @@ +ALTER TYPE resource_type ADD VALUE IF NOT EXISTS 'ai_seat'; diff --git a/coderd/database/models.go b/coderd/database/models.go index 9b07ae580a..2ac42f9277 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -3027,6 +3027,7 @@ const ( ResourceTypeWorkspaceApp ResourceType = "workspace_app" ResourceTypePrebuildsSettings ResourceType = "prebuilds_settings" ResourceTypeTask ResourceType = "task" + ResourceTypeAiSeat ResourceType = "ai_seat" ) func (e *ResourceType) Scan(src interface{}) error { @@ -3091,7 +3092,8 @@ func (e ResourceType) Valid() bool { ResourceTypeWorkspaceAgent, ResourceTypeWorkspaceApp, ResourceTypePrebuildsSettings, - ResourceTypeTask: + ResourceTypeTask, + ResourceTypeAiSeat: return true } return false @@ -3125,6 +3127,7 @@ func AllResourceTypeValues() []ResourceType { ResourceTypeWorkspaceApp, ResourceTypePrebuildsSettings, ResourceTypeTask, + ResourceTypeAiSeat, } } diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c252191483..dfb368c03d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -835,7 +835,8 @@ 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 + // Returns true if a new rows was inserted, false otherwise. + UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, 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/querier_test.go b/coderd/database/querier_test.go index df1e278b83..988e550c4d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -9441,3 +9441,42 @@ func TestGetWorkspaceBuildMetricsByResourceID(t *testing.T) { require.Equal(t, "success", row.WorstStatus) }) } + +// TestUpsertAISeats verifies 'UpsertAISeatState' only returns true when a new +// row is inserted. +func TestUpsertAISeats(t *testing.T) { + t.Parallel() + + sqlDB := testSQLDB(t) + err := migrations.Up(sqlDB) + require.NoError(t, err) + db := database.New(sqlDB) + ctx := testutil.Context(t, testutil.WaitShort) + + now := dbtime.Now() + + user := dbgen.User(t, db, database.User{}) + newRow, err := db.UpsertAISeatState(ctx, database.UpsertAISeatStateParams{ + UserID: user.ID, + FirstUsedAt: now.Add(time.Hour * -24), + LastEventType: database.AiSeatUsageReasonTask, + }) + require.NoError(t, err) + require.True(t, newRow) + + alreadyExists, err := db.UpsertAISeatState(ctx, database.UpsertAISeatStateParams{ + UserID: user.ID, + FirstUsedAt: now.Add(time.Hour * -23), + LastEventType: database.AiSeatUsageReasonTask, + }) + require.NoError(t, err) + require.False(t, alreadyExists) + + alreadyExists, err = db.UpsertAISeatState(ctx, database.UpsertAISeatStateParams{ + UserID: user.ID, + FirstUsedAt: now, + LastEventType: database.AiSeatUsageReasonTask, + }) + require.NoError(t, err) + require.False(t, alreadyExists) +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 79283ac791..7376d01004 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1220,7 +1220,7 @@ func (q *sqlQuerier) GetActiveAISeatCount(ctx context.Context) (int64, error) { return count, err } -const upsertAISeatState = `-- name: UpsertAISeatState :exec +const upsertAISeatState = `-- name: UpsertAISeatState :one INSERT INTO ai_seat_state ( user_id, first_used_at, @@ -1237,6 +1237,9 @@ SET last_event_type = EXCLUDED.last_event_type, last_event_description = EXCLUDED.last_event_description, updated_at = EXCLUDED.updated_at +RETURNING + -- Postgres vodoo to know if a row was inserted. + (xmax = 0)::boolean AS is_new ` type UpsertAISeatStateParams struct { @@ -1246,14 +1249,17 @@ type UpsertAISeatStateParams struct { 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, +// Returns true if a new rows was inserted, false otherwise. +func (q *sqlQuerier) UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error) { + row := q.db.QueryRowContext(ctx, upsertAISeatState, arg.UserID, arg.FirstUsedAt, arg.LastEventType, arg.LastEventDescription, ) - return err + var is_new bool + err := row.Scan(&is_new) + return is_new, err } const deleteAPIKeyByID = `-- name: DeleteAPIKeyByID :exec diff --git a/coderd/database/queries/aiseats.sql b/coderd/database/queries/aiseats.sql index 628e2fffd6..39e1d76b19 100644 --- a/coderd/database/queries/aiseats.sql +++ b/coderd/database/queries/aiseats.sql @@ -1,4 +1,5 @@ --- name: UpsertAISeatState :exec +-- name: UpsertAISeatState :one +-- Returns true if a new rows was inserted, false otherwise. INSERT INTO ai_seat_state ( user_id, first_used_at, @@ -15,7 +16,9 @@ SET last_event_type = EXCLUDED.last_event_type, last_event_description = EXCLUDED.last_event_description, updated_at = EXCLUDED.updated_at -; +RETURNING + -- Postgres vodoo to know if a row was inserted. + (xmax = 0)::boolean AS is_new; -- name: GetActiveAISeatCount :one SELECT diff --git a/codersdk/audit.go b/codersdk/audit.go index 0b2eca7d79..5018982c6c 100644 --- a/codersdk/audit.go +++ b/codersdk/audit.go @@ -45,6 +45,7 @@ const ( // connection log. ResourceTypeWorkspaceApp ResourceType = "workspace_app" ResourceTypeTask ResourceType = "task" + ResourceTypeAISeat ResourceType = "ai_seat" ) func (r ResourceType) FriendlyString() string { @@ -103,6 +104,8 @@ func (r ResourceType) FriendlyString() string { return "workspace app" case ResourceTypeTask: return "task" + case ResourceTypeAISeat: + return "ai seat" default: return "unknown" } diff --git a/docs/admin/security/audit-logs.md b/docs/admin/security/audit-logs.md index 5d45031086..aed4b96291 100644 --- a/docs/admin/security/audit-logs.md +++ b/docs/admin/security/audit-logs.md @@ -16,6 +16,7 @@ We track the following resources: | Resource | | | |-----------------------------------------------------------------|----------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | APIKey
login, logout, register, create, write, delete | |
FieldTracked
allow_listfalse
created_attrue
expires_attrue
hashed_secretfalse
idfalse
ip_addressfalse
last_usedtrue
lifetime_secondsfalse
login_typefalse
scopesfalse
token_namefalse
updated_atfalse
user_idtrue
| +| AiSeatState
create | |
FieldTracked
first_used_attrue
last_event_descriptiontrue
last_event_typetrue
last_used_atfalse
updated_atfalse
user_idtrue
| | AuditOAuthConvertState
| |
FieldTracked
created_attrue
expires_attrue
from_login_typetrue
to_login_typetrue
user_idtrue
| | Group
create, write, delete | |
FieldTracked
avatar_urltrue
display_nametrue
idtrue
memberstrue
nametrue
organization_idfalse
quota_allowancetrue
sourcefalse
| | AuditableOrganizationMember
| |
FieldTracked
created_attrue
organization_idfalse
rolestrue
updated_attrue
user_idtrue
usernametrue
| diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 3b563d3bb9..80a0ccec90 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7557,9 +7557,9 @@ Only certain features set these fields: - FeatureManagedAgentLimit| #### Enumerated Values -| Value(s) | -|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `api_key`, `convert_login`, `custom_role`, `git_ssh_key`, `group`, `health_settings`, `idp_sync_settings_group`, `idp_sync_settings_organization`, `idp_sync_settings_role`, `license`, `notification_template`, `notifications_settings`, `oauth2_provider_app`, `oauth2_provider_app_secret`, `organization`, `organization_member`, `prebuilds_settings`, `task`, `template`, `template_version`, `user`, `workspace`, `workspace_agent`, `workspace_app`, `workspace_build`, `workspace_proxy` | +| Value(s) | +|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `ai_seat`, `api_key`, `convert_login`, `custom_role`, `git_ssh_key`, `group`, `health_settings`, `idp_sync_settings_group`, `idp_sync_settings_organization`, `idp_sync_settings_role`, `license`, `notification_template`, `notifications_settings`, `oauth2_provider_app`, `oauth2_provider_app_secret`, `organization`, `organization_member`, `prebuilds_settings`, `task`, `template`, `template_version`, `user`, `workspace`, `workspace_agent`, `workspace_app`, `workspace_build`, `workspace_proxy` | ## codersdk.Response diff --git a/enterprise/aibridgedserver/aibridgedserver_test.go b/enterprise/aibridgedserver/aibridgedserver_test.go index b195829534..e38890a5fc 100644 --- a/enterprise/aibridgedserver/aibridgedserver_test.go +++ b/enterprise/aibridgedserver/aibridgedserver_test.go @@ -24,7 +24,7 @@ import ( "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogjson" - "github.com/coder/coder/v2/coderd/aiseats" + agplaiseats "github.com/coder/coder/v2/coderd/aiseats" "github.com/coder/coder/v2/coderd/apikey" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" @@ -177,7 +177,7 @@ func TestAuthorization(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, aiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) require.NotNil(t, srv) @@ -269,7 +269,7 @@ func TestGetMCPServerConfigs(t *testing.T) { accessURL := "https://my-cool-deployment.com" srv, err := aibridgedserver.NewServer(t.Context(), db, logger, accessURL, codersdk.AIBridgeConfig{ InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection), - }, tc.externalAuthConfigs, tc.experiments, aiseats.Noop{}) + }, tc.externalAuthConfigs, tc.experiments, agplaiseats.Noop{}) require.NoError(t, err) require.NotNil(t, srv) @@ -319,7 +319,7 @@ func TestGetMCPServerAccessTokensBatch(t *testing.T) { { ID: "3", }, - }, requiredExperiments, aiseats.Noop{}) + }, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) require.NotNil(t, srv) @@ -1015,7 +1015,7 @@ func testRecordMethod[Req any, Resp any]( } ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, aiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) resp, err := callMethod(srv, ctx, tc.request) @@ -1310,7 +1310,7 @@ func TestStructuredLogging(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{ StructuredLogging: serpent.Bool(tc.structuredLogging), - }, nil, requiredExperiments, aiseats.Noop{}) + }, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) err = tc.recordFn(srv, ctx, interceptionID) @@ -1352,7 +1352,7 @@ func TestInferredThreadsByToolCalls(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, aiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) aID := uuid.New() diff --git a/enterprise/aiseats/tracker.go b/enterprise/aiseats/tracker.go index d69dc0491c..30cd8abfb5 100644 --- a/enterprise/aiseats/tracker.go +++ b/enterprise/aiseats/tracker.go @@ -3,18 +3,20 @@ package aiseats import ( "context" "sync" + "sync/atomic" "time" "github.com/google/uuid" "cdr.dev/slog/v3" agplaiseats "github.com/coder/coder/v2/coderd/aiseats" + "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/database" "github.com/coder/quartz" ) type store interface { - UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) error + UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) } // throttleInterval is the minimum time between DB writes for the same user. This @@ -33,19 +35,20 @@ const ( // SeatTracker records current AI seat state for users. type SeatTracker struct { - db store - logger slog.Logger - clock quartz.Clock + db store + logger slog.Logger + clock quartz.Clock + auditor *atomic.Pointer[audit.Auditor] mu sync.RWMutex retryAfter map[uuid.UUID]time.Time } -func New(db store, logger slog.Logger, clock quartz.Clock) *SeatTracker { +func New(db store, logger slog.Logger, clock quartz.Clock, auditor *atomic.Pointer[audit.Auditor]) *SeatTracker { if clock == nil { clock = quartz.NewReal() } - return &SeatTracker{db: db, logger: logger, clock: clock, retryAfter: make(map[uuid.UUID]time.Time)} + return &SeatTracker{db: db, logger: logger, clock: clock, auditor: auditor, retryAfter: make(map[uuid.UUID]time.Time)} } // skipRecord returns true when the user is still in the retry cooldown @@ -75,7 +78,7 @@ func (t *SeatTracker) RecordUsage(ctx context.Context, userID uuid.UUID, reason return } - err := t.db.UpsertAISeatState(ctx, database.UpsertAISeatStateParams{ + isNew, err := t.db.UpsertAISeatState(ctx, database.UpsertAISeatStateParams{ UserID: userID, FirstUsedAt: now, LastEventType: reason.EventType, @@ -88,4 +91,26 @@ func (t *SeatTracker) RecordUsage(ctx context.Context, userID uuid.UUID, reason } t.recordThrottle(userID, now, throttleInterval) + if isNew && t.auditor != nil { + // Record an audit log for the first time a user uses an AI seat. + auditor := t.auditor.Load() + if auditor == nil || *auditor == nil { + return + } + audit.BackgroundAudit[database.AiSeatState](ctx, &audit.BackgroundAuditParams[database.AiSeatState]{ + Audit: *auditor, + Log: t.logger, + UserID: userID, + Time: now, + Action: database.AuditActionCreate, + New: database.AiSeatState{ + UserID: userID, + FirstUsedAt: now, + LastUsedAt: now, + LastEventType: reason.EventType, + LastEventDescription: reason.Description, + UpdatedAt: now, + }, + }) + } } diff --git a/enterprise/aiseats/tracker_test.go b/enterprise/aiseats/tracker_test.go index 574e80fcbb..cbebd7a077 100644 --- a/enterprise/aiseats/tracker_test.go +++ b/enterprise/aiseats/tracker_test.go @@ -1,12 +1,14 @@ package aiseats_test import ( + "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" agplaiseats "github.com/coder/coder/v2/coderd/aiseats" + "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" @@ -25,7 +27,7 @@ func TestSeatTrackerDB(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) clock := quartz.NewMock(t) - tracker := enterpriseaiseats.New(db, testutil.Logger(t), clock) + tracker := enterpriseaiseats.New(db, testutil.Logger(t), clock, nil) user := dbgen.User(t, db, database.User{Status: database.UserStatusActive}) tracker.RecordUsage(ctx, user.ID, agplaiseats.ReasonAIBridge("active user event")) @@ -40,7 +42,7 @@ func TestSeatTrackerDB(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - tracker := enterpriseaiseats.New(db, testutil.Logger(t), quartz.NewMock(t)) + tracker := enterpriseaiseats.New(db, testutil.Logger(t), quartz.NewMock(t), nil) dormantUser := dbgen.User(t, db, database.User{Status: database.UserStatusDormant}) tracker.RecordUsage(ctx, dormantUser.ID, agplaiseats.ReasonTask("dormant user event")) @@ -58,7 +60,12 @@ func TestSeatTrackerDB(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) - tracker := enterpriseaiseats.New(db, testutil.Logger(t), quartz.NewMock(t)) + a := audit.NewMock() + var aI audit.Auditor = a + var al atomic.Pointer[audit.Auditor] + al.Store(&aI) + + tracker := enterpriseaiseats.New(db, testutil.Logger(t), quartz.NewMock(t), &al) user := dbgen.User(t, db, database.User{Status: database.UserStatusActive}) tracker.RecordUsage(ctx, user.ID, agplaiseats.ReasonAIBridge("status transition")) @@ -90,5 +97,8 @@ func TestSeatTrackerDB(t *testing.T) { count, err = db.GetActiveAISeatCount(ctx) require.NoError(t, err) require.EqualValues(t, 1, count) + + require.Len(t, a.AuditLogs(), 1) + require.Equal(t, database.ResourceTypeAiSeat, a.AuditLogs()[0].ResourceType) }) } diff --git a/enterprise/audit/table.go b/enterprise/audit/table.go index a226b5d630..d556e47d44 100644 --- a/enterprise/audit/table.go +++ b/enterprise/audit/table.go @@ -28,6 +28,7 @@ var AuditActionMap = map[string][]codersdk.AuditAction{ "APIKey": {codersdk.AuditActionLogin, codersdk.AuditActionLogout, codersdk.AuditActionRegister, codersdk.AuditActionCreate, codersdk.AuditActionWrite, codersdk.AuditActionDelete}, "License": {codersdk.AuditActionCreate, codersdk.AuditActionDelete}, "Task": {codersdk.AuditActionCreate, codersdk.AuditActionWrite, codersdk.AuditActionDelete}, + "AiSeatState": {codersdk.AuditActionCreate}, } type Action string @@ -351,6 +352,17 @@ var auditableResourcesTypes = map[any]map[string]Action{ "field": ActionTrack, "mapping": ActionTrack, }, + &database.AiSeatState{}: { + "user_id": ActionTrack, + "first_used_at": ActionTrack, + "last_event_type": ActionTrack, + "last_event_description": ActionTrack, + + // Since the audit log only fires on the first event, these fields will always + // match "first_used_at". + "last_used_at": ActionIgnore, + "updated_at": ActionIgnore, + }, &database.TaskTable{}: { "id": ActionTrack, "organization_id": ActionIgnore, // Never changes. diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index bfefc89404..3c29f6454f 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -218,9 +218,8 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }, }) - api.aiSeatTracker = aiseats.New(options.Database, api.Logger.Named("aiseats"), quartz.NewReal()) - api.AGPL = coderd.New(options.Options) + api.aiSeatTracker = aiseats.New(options.Database, api.Logger.Named("aiseats"), quartz.NewReal(), &api.AGPL.Auditor) api.AGPL.AISeatTracker = api.aiSeatTracker defer func() { if err != nil { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 0bc0bd3aa6..f12a8cc36e 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -5155,6 +5155,7 @@ export interface ResolveAutostartResponse { // From codersdk/audit.go export type ResourceType = + | "ai_seat" | "api_key" | "convert_login" | "custom_role" @@ -5183,6 +5184,7 @@ export type ResourceType = | "workspace_proxy"; export const ResourceTypes: ResourceType[] = [ + "ai_seat", "api_key", "convert_login", "custom_role",