coder · ssncferreira · Jul 28, 2026 · Jul 24, 2026 · Jul 27, 2026 · Jul 27, 2026
diff --git a/cli/server.go b/cli/server.go
@@ -65,6 +65,7 @@ import (
 	"github.com/coder/coder/v2/cli/config"
 	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/aibridged"
+	"github.com/coder/coder/v2/coderd/aibridgedserver"
 	"github.com/coder/coder/v2/coderd/authlink"
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/cryptokeys"
@@ -1183,6 +1184,16 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				// https://linear.app/codercom/issue/AIGOV-447/remove-legacy-ai-gateway-metric-aliases
 				aibridgeReg := prometheusmetrics.NewMetricAliasRegisterer(coderAPI.PrometheusRegistry, aibridgemetrics.PrometheusMetricPrefix, "coder_aibridged_")
 				aibridgeMetrics := aibridge.NewMetrics(aibridgeReg)
+				costControlReg := prometheus.WrapRegistererWithPrefix("coder_ai_gateway_", coderAPI.PrometheusRegistry)
+				coderAPI.AIGatewayServerMetrics = aibridgedserver.NewMetrics(costControlReg)
+				if vals.Prometheus.Enable {
+					budgetPeriod := codersdk.NewAIBudgetPeriodFromString(vals.AI.BridgeConfig.BudgetPeriod)
+					closeBlockedUsersFunc := coderAPI.AIGatewayServerMetrics.StartBlockedUsersCollector(
+						ctx, logger.Named("aigateway_cost_control_metrics"), quartz.NewReal(),
+						coderAPI.Database, budgetPeriod, 0,
+					)
+					defer closeBlockedUsersFunc()
+				}
 				var unsubscribeProviderReload func()
 				aibridgeDaemon, unsubscribeProviderReload, err = newAIBridgeDaemon(coderAPI, vals.AI.BridgeConfig, aibridgeReg, aibridgeMetrics)
 				if err != nil {

diff --git a/coderd/aibridged.go b/coderd/aibridged.go
@@ -76,6 +76,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai
 		Experiments:         api.Experiments,
 		Logger:              api.Logger.Named("aibridgedserver"),
 		Clock:               api.Clock,
+		Metrics:             api.AIGatewayServerMetrics,
 	})
 	if err != nil {
 		return nil, err

diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go
@@ -124,6 +124,8 @@ type Server struct {
 	budgetPeriod  codersdk.AIBudgetPeriod
 	clock         quartz.Clock
 	notifEnqueuer notifications.Enqueuer
+	// metrics records cost-control metrics. May be nil.
+	metrics *Metrics
 }
 
 // Options carries the dependencies required to construct an aibridged Server.
@@ -140,8 +142,9 @@ type Options struct {
 	ExternalAuthConfigs []*externalauth.Config
 	Experiments         codersdk.Experiments
 
-	Logger slog.Logger
-	Clock  quartz.Clock
+	Logger  slog.Logger
+	Clock   quartz.Clock
+	Metrics *Metrics
 }
 
 func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) {
@@ -173,6 +176,7 @@ func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) {
 		budgetPeriod:        codersdk.NewAIBudgetPeriodFromString(opts.GatewayCfg.BudgetPeriod),
 		clock:               opts.Clock,
 		notifEnqueuer:       enqueuer,
+		metrics:             opts.Metrics,
 	}
 
 	if opts.GatewayCfg.InjectCoderMCPTools {
@@ -811,10 +815,25 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest
 // IsBudgetExceeded reports whether the user's AI spend has reached their
 // effective limit over [periodStart, now], where periodStart is the start of
 // the current deployment-configured budget period.
-func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (*proto.IsBudgetExceededResponse, error) {
+func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (resp *proto.IsBudgetExceededResponse, retErr error) {
 	//nolint:gocritic // AIBridged has specific authz rules.
 	ctx = dbauthz.AsAIBridged(ctx)
 
+	start := s.clock.Now()
+	defer func() {
+		if s.metrics == nil {
+			return
+		}
+		outcome := "allowed"
+		switch {
+		case retErr != nil:
+			outcome = "error"
+		case resp != nil && resp.Exceeded:
+			outcome = "blocked"
+		}
+		s.metrics.EnforcementDuration.WithLabelValues(outcome).Observe(s.clock.Since(start).Seconds())
+	}()
+
 	userID, err := uuid.Parse(in.GetUserId())
 	if err != nil {
 		return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err)
@@ -829,17 +848,22 @@ func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceede
 	if err != nil {
 		return nil, err
 	}
+	if userBudget.Exceeded && s.metrics != nil {
+		s.metrics.BlockedRequests.WithLabelValues(userBudget.GroupID.String()).Inc()
+	}
 	return &proto.IsBudgetExceededResponse{
 		Exceeded:         userBudget.Exceeded,
 		SpendLimitMicros: userBudget.SpendLimitMicros,
 	}, nil
 }
 
 // userAIBudget is a snapshot of a user's AI budget status. SpendLimitMicros
-// is nil when no budget is configured for the user (unlimited).
+// is nil when no budget is configured for the user (unlimited). GroupID is the
+// effective group the limit resolved to, set only when a limit applies.
 type userAIBudget struct {
 	Exceeded         bool
 	SpendLimitMicros *int64
+	GroupID          uuid.UUID
 }
 
 // checkUserAIBudget evaluates the user's AI budget status aggregated over
@@ -894,6 +918,7 @@ func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, period
 	return userAIBudget{
 		Exceeded:         exceeded,
 		SpendLimitMicros: ptr.Ref(effectiveGroup.Limit.SpendLimitMicros),
+		GroupID:          effectiveGroup.GroupID,
 	}, nil
 }