opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: add audit log entry when ai seat is consumed (#22683)

Browse Source 
When an ai seat is consumed, an audit log entry is made. This only happens the first time a seat is used.
This commit is contained in:
Steven Masley
2026-03-16 15:30:25 -05:00
committed by GitHub
parent 6972d073a2
commit 93b9d70a9b
26 changed files with 167 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/apidoc/docs.go
Generated
+4 -2
View File
@@ -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": {
coderd/apidoc/swagger.json
Generated
+4 -2
View File
@@ -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": {
coderd/audit/diff.go
+2 -1
View File
@@ -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
coderd/audit/request.go
+8
View File
@@ -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))
}
coderd/coderd.go
+1 -1
View File
@@ -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,
coderd/database/dbauthz/dbauthz.go
+2 -2
View File
@@ -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)
}
coderd/database/dbauthz/dbauthz_test.go
+1 -1
View File
@@ -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) {
coderd/database/dbmetrics/querymetrics.go
+3 -3
View File
@@ -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 {
coderd/database/dbmock/dbmock.go
+4 -3
View File
@@ -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.
coderd/database/dump.sql
Generated
+2 -1
View File
@@ -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 (
coderd/database/migrations/000440_ai_seat_audit.down.sql
+1
View File
@@ -0,0 +1 @@
-- resource_type enum values cannot be removed safely; no-op.
coderd/database/migrations/000440_ai_seat_audit.up.sql
+1
View File
@@ -0,0 +1 @@
ALTER TYPE resource_type ADD VALUE IF NOT EXISTS 'ai_seat';
coderd/database/models.go
+4 -1
View File
@@ -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,
}
}
coderd/database/querier.go
+2 -1
View File
@@ -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
coderd/database/querier_test.go
+39
View File
@@ -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)
}
coderd/database/queries.sql.go
+10 -4
View File
@@ -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
coderd/database/queries/aiseats.sql
+5 -2
View File
@@ -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
codersdk/audit.go
+3
View File
@@ -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"
}
docs/admin/security/audit-logs.md
+1
View File
@@ -16,6 +16,7 @@ We track the following resources:
| <b>Resource<b> | | |
|-----------------------------------------------------------------|----------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| APIKey<br><i>login, logout, register, create, write, delete</i> | <table><thead><tr><th>Field</th><th>Tracked</th></tr></thead><tbody> | <tr><td>allow_list</td><td>false</td></tr><tr><td>created_at</td><td>true</td></tr><tr><td>expires_at</td><td>true</td></tr><tr><td>hashed_secret</td><td>false</td></tr><tr><td>id</td><td>false</td></tr><tr><td>ip_address</td><td>false</td></tr><tr><td>last_used</td><td>true</td></tr><tr><td>lifetime_seconds</td><td>false</td></tr><tr><td>login_type</td><td>false</td></tr><tr><td>scopes</td><td>false</td></tr><tr><td>token_name</td><td>false</td></tr><tr><td>updated_at</td><td>false</td></tr><tr><td>user_id</td><td>true</td></tr></tbody></table> |
| AiSeatState<br><i>create</i> | <table><thead><tr><th>Field</th><th>Tracked</th></tr></thead><tbody> | <tr><td>first_used_at</td><td>true</td></tr><tr><td>last_event_description</td><td>true</td></tr><tr><td>last_event_type</td><td>true</td></tr><tr><td>last_used_at</td><td>false</td></tr><tr><td>updated_at</td><td>false</td></tr><tr><td>user_id</td><td>true</td></tr></tbody></table> |
| AuditOAuthConvertState<br><i></i> | <table><thead><tr><th>Field</th><th>Tracked</th></tr></thead><tbody> | <tr><td>created_at</td><td>true</td></tr><tr><td>expires_at</td><td>true</td></tr><tr><td>from_login_type</td><td>true</td></tr><tr><td>to_login_type</td><td>true</td></tr><tr><td>user_id</td><td>true</td></tr></tbody></table> |
| Group<br><i>create, write, delete</i> | <table><thead><tr><th>Field</th><th>Tracked</th></tr></thead><tbody> | <tr><td>avatar_url</td><td>true</td></tr><tr><td>display_name</td><td>true</td></tr><tr><td>id</td><td>true</td></tr><tr><td>members</td><td>true</td></tr><tr><td>name</td><td>true</td></tr><tr><td>organization_id</td><td>false</td></tr><tr><td>quota_allowance</td><td>true</td></tr><tr><td>source</td><td>false</td></tr></tbody></table> |
| AuditableOrganizationMember<br><i></i> | <table><thead><tr><th>Field</th><th>Tracked</th></tr></thead><tbody> | <tr><td>created_at</td><td>true</td></tr><tr><td>organization_id</td><td>false</td></tr><tr><td>roles</td><td>true</td></tr><tr><td>updated_at</td><td>true</td></tr><tr><td>user_id</td><td>true</td></tr><tr><td>username</td><td>true</td></tr></tbody></table> |
docs/reference/api/schemas.md
Generated
+3 -3
View File
@@ -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
enterprise/aibridgedserver/aibridgedserver_test.go
+7 -7
View File
@@ -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()
enterprise/aiseats/tracker.go
+32 -7
View File
@@ -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,
},
})
}
}
enterprise/aiseats/tracker_test.go
+13 -3
View File
@@ -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)
})
}
enterprise/audit/table.go
+12
View File
@@ -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.
enterprise/coderd/coderd.go
+1 -2
View File
@@ -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 {
site/src/api/typesGenerated.ts
Generated
+2
View File
@@ -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",