mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add Prometheus metrics for AI Governance cost control (#27490)
## Description Adds Prometheus metrics for AI budget cost control, emitted by the aibridged server under the `cost_control` subsystem (full names are prefixed `coder_ai_gateway_`). - `blocked_requests_total` (counter) — labels: `group_id` - `blocked_users` (gauge) — labels: `group_id` - `unpriced_requests_total` (counter) — labels: `provider`, `model` - `enforcement_duration_seconds` (histogram) — labels: `outcome` ## Changes - Add `GetOverBudgetUsersPerGroup` query (plus dbauthz/dbmetrics/dbmock wiring) to count over-budget users per effective group. - Add a background collector that refreshes the `blocked_users` gauge on an interval, started only when Prometheus is enabled. - Wire `Metrics` through the aibridged server, coderd API, `cli/server.go`, and the enterprise AI gateway handler; recording is nil-safe when metrics are unset. Closes https://linear.app/codercom/issue/AIGOV-296/add-prometheus-metrics-for-cost-control > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
@@ -4321,6 +4321,14 @@ func (q *querier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg da
|
||||
return q.db.GetOrganizationsWithPrebuildStatus(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetOverBudgetUsersPerGroup(ctx context.Context, periodStart time.Time) ([]database.GetOverBudgetUsersPerGroupRow, error) {
|
||||
// Aggregates over-budget user counts per group for cost-control metrics.
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceGroup.All()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetOverBudgetUsersPerGroup(ctx, periodStart)
|
||||
}
|
||||
|
||||
func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
|
||||
version, err := q.db.GetTemplateVersionByJobID(ctx, jobID)
|
||||
if err != nil {
|
||||
|
||||
@@ -7073,6 +7073,13 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
check.Args(arg).Asserts(user, policy.ActionRead).Returns(row)
|
||||
}))
|
||||
|
||||
s.Run("GetOverBudgetUsersPerGroup", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
periodStart := time.Now().UTC().Truncate(24 * time.Hour)
|
||||
dbm.EXPECT().GetOverBudgetUsersPerGroup(gomock.Any(), periodStart).
|
||||
Return([]database.GetOverBudgetUsersPerGroupRow{}, nil).AnyTimes()
|
||||
check.Args(periodStart).Asserts(rbac.ResourceGroup.All(), policy.ActionRead)
|
||||
}))
|
||||
|
||||
s.Run("IncrementUserAIDailySpend", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.IncrementUserAIDailySpendParams{
|
||||
UserID: uuid.New(),
|
||||
|
||||
Reference in New Issue
Block a user