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

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
+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": {
+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": {
+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
+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))
}
+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,
+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)
}
+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) {
+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 {
+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.
+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 (
@@ -0,0 +1 @@
-- resource_type enum values cannot be removed safely; no-op.
@@ -0,0 +1 @@
ALTER TYPE resource_type ADD VALUE IF NOT EXISTS 'ai_seat';
+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,
}
}
+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
+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)
}
+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
+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