diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index e2b913dcd5b..d8d02c13146 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" "github.com/coder/serpent" ) @@ -53,7 +54,7 @@ func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provide // (providers, outcomes) the embedded reloader would observe. func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) { t.Helper() - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", cfg, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) if err != nil { return nil, nil, err } diff --git a/coderd/aibridged.go b/coderd/aibridged.go index 0749bb2cec7..3784545f7c7 100644 --- a/coderd/aibridged.go +++ b/coderd/aibridged.go @@ -66,7 +66,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai mux := drpcmux.New() srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Pubsub, api.Logger.Named("aibridgedserver"), - api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker) + api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker, api.Clock) if err != nil { return nil, err } diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go index 7c9fef2cd54..0f8c099eaaa 100644 --- a/coderd/aibridged/http.go +++ b/coderd/aibridged/http.go @@ -7,16 +7,12 @@ import ( "github.com/google/uuid" "golang.org/x/xerrors" - "google.golang.org/protobuf/types/known/timestamppb" "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge" "github.com/coder/coder/v2/aibridge/recorder" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" - "github.com/coder/coder/v2/coderd/aibridge/budget" "github.com/coder/coder/v2/coderd/aibridged/proto" - "github.com/coder/coder/v2/coderd/database/dbtime" - "github.com/coder/coder/v2/codersdk" ) var _ http.Handler = &Server{} @@ -149,15 +145,8 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { } logger = logger.With(slog.F("user_id", id)) - periodWindow, err := budget.CurrentPeriod(dbtime.Now(), codersdk.AIBudgetPeriodMonth) - if err != nil { - logger.Warn(ctx, "compute AI budget period", slog.Error(err)) - http.Error(rw, ErrBudgetCheck.Error(), http.StatusInternalServerError) - return - } budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ - UserId: id.String(), - PeriodStart: timestamppb.New(periodWindow.Start), + UserId: id.String(), }) if err != nil { logger.Warn(ctx, "user AI budget check failed", slog.Error(err)) diff --git a/coderd/aibridged/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index fe68d6ee205..e029c0c3931 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1305,8 +1305,6 @@ type IsBudgetExceededRequest struct { unknownFields protoimpl.UnknownFields UserId string `protobuf:"bytes,1,opt,name=user_id,json=userId,proto3" json:"user_id,omitempty"` // UUID - // The spend aggregation window is [period_start, now]. - PeriodStart *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=period_start,json=periodStart,proto3" json:"period_start,omitempty"` } func (x *IsBudgetExceededRequest) Reset() { @@ -1348,13 +1346,6 @@ func (x *IsBudgetExceededRequest) GetUserId() string { return "" } -func (x *IsBudgetExceededRequest) GetPeriodStart() *timestamppb.Timestamp { - if x != nil { - return x.PeriodStart - } - return nil -} - type IsBudgetExceededResponse struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache @@ -2026,135 +2017,131 @@ var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ 0x0a, 0x0a, 0x61, 0x70, 0x69, 0x5f, 0x6b, 0x65, 0x79, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x61, 0x70, 0x69, 0x4b, 0x65, 0x79, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x75, 0x73, 0x65, 0x72, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, - 0x75, 0x73, 0x65, 0x72, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x71, 0x0a, 0x17, 0x49, 0x73, 0x42, 0x75, + 0x75, 0x73, 0x65, 0x72, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x32, 0x0a, 0x17, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x17, 0x0a, 0x07, 0x75, 0x73, 0x65, 0x72, 0x5f, 0x69, 0x64, 0x18, 0x01, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x75, 0x73, 0x65, 0x72, 0x49, 0x64, 0x12, 0x3d, 0x0a, 0x0c, - 0x70, 0x65, 0x72, 0x69, 0x6f, 0x64, 0x5f, 0x73, 0x74, 0x61, 0x72, 0x74, 0x18, 0x02, 0x20, 0x01, - 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x0b, - 0x70, 0x65, 0x72, 0x69, 0x6f, 0x64, 0x53, 0x74, 0x61, 0x72, 0x74, 0x22, 0x80, 0x01, 0x0a, 0x18, - 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, - 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x1a, 0x0a, 0x08, 0x65, 0x78, 0x63, 0x65, - 0x65, 0x64, 0x65, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x08, 0x65, 0x78, 0x63, 0x65, - 0x65, 0x64, 0x65, 0x64, 0x12, 0x31, 0x0a, 0x12, 0x73, 0x70, 0x65, 0x6e, 0x64, 0x5f, 0x6c, 0x69, - 0x6d, 0x69, 0x74, 0x5f, 0x6d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, - 0x48, 0x00, 0x52, 0x10, 0x73, 0x70, 0x65, 0x6e, 0x64, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x4d, 0x69, - 0x63, 0x72, 0x6f, 0x73, 0x88, 0x01, 0x01, 0x42, 0x15, 0x0a, 0x13, 0x5f, 0x73, 0x70, 0x65, 0x6e, - 0x64, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x5f, 0x6d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x22, 0x17, - 0x0a, 0x15, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x49, 0x0a, 0x16, 0x47, 0x65, 0x74, 0x41, 0x49, - 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x12, 0x2f, 0x0a, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x18, 0x01, - 0x20, 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, - 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x52, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, - 0x72, 0x73, 0x22, 0x19, 0x0a, 0x17, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, - 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x1a, 0x0a, - 0x18, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, - 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0xb5, 0x01, 0x0a, 0x0a, 0x41, 0x49, - 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, - 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x12, 0x0a, 0x04, - 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, - 0x12, 0x18, 0x0a, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, - 0x08, 0x52, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x19, 0x0a, 0x08, 0x62, 0x61, - 0x73, 0x65, 0x5f, 0x75, 0x72, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, 0x61, - 0x73, 0x65, 0x55, 0x72, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x18, 0x05, 0x20, - 0x03, 0x28, 0x09, 0x52, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x12, 0x36, 0x0a, 0x07, 0x62, 0x65, 0x64, - 0x72, 0x6f, 0x63, 0x6b, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, 0x6e, - 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x52, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, 0x63, - 0x6b, 0x22, 0xf6, 0x01, 0x0a, 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, - 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x12, 0x16, 0x0a, 0x06, 0x72, - 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, 0x67, - 0x69, 0x6f, 0x6e, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, - 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x4b, - 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x11, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, 0x79, - 0x5f, 0x73, 0x65, 0x63, 0x72, 0x65, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0f, 0x61, - 0x63, 0x63, 0x65, 0x73, 0x73, 0x4b, 0x65, 0x79, 0x53, 0x65, 0x63, 0x72, 0x65, 0x74, 0x12, 0x14, - 0x0a, 0x05, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6d, - 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x28, 0x0a, 0x10, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x5f, 0x66, 0x61, - 0x73, 0x74, 0x5f, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0e, - 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x46, 0x61, 0x73, 0x74, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x19, - 0x0a, 0x08, 0x72, 0x6f, 0x6c, 0x65, 0x5f, 0x61, 0x72, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, - 0x52, 0x07, 0x72, 0x6f, 0x6c, 0x65, 0x41, 0x72, 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x65, 0x78, 0x74, - 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x5f, 0x69, 0x64, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, - 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x49, 0x64, 0x32, 0xa9, 0x04, 0x0a, 0x08, 0x52, - 0x65, 0x63, 0x6f, 0x72, 0x64, 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, 0x72, - 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, - 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, - 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, - 0x73, 0x65, 0x12, 0x68, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, - 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, - 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, - 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, - 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, - 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, - 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, - 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, - 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x12, 0x56, 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, - 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, - 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x75, 0x73, 0x65, 0x72, 0x49, 0x64, 0x22, 0x80, 0x01, 0x0a, + 0x18, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, + 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x1a, 0x0a, 0x08, 0x65, 0x78, 0x63, + 0x65, 0x65, 0x64, 0x65, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x08, 0x65, 0x78, 0x63, + 0x65, 0x65, 0x64, 0x65, 0x64, 0x12, 0x31, 0x0a, 0x12, 0x73, 0x70, 0x65, 0x6e, 0x64, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x5f, 0x6d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x03, 0x48, 0x00, 0x52, 0x10, 0x73, 0x70, 0x65, 0x6e, 0x64, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x4d, + 0x69, 0x63, 0x72, 0x6f, 0x73, 0x88, 0x01, 0x01, 0x42, 0x15, 0x0a, 0x13, 0x5f, 0x73, 0x70, 0x65, + 0x6e, 0x64, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x5f, 0x6d, 0x69, 0x63, 0x72, 0x6f, 0x73, 0x22, + 0x17, 0x0a, 0x15, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, + 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x49, 0x0a, 0x16, 0x47, 0x65, 0x74, 0x41, + 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x2f, 0x0a, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x18, + 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, + 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x52, 0x09, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, + 0x65, 0x72, 0x73, 0x22, 0x19, 0x0a, 0x17, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, + 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0x1a, + 0x0a, 0x18, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0xb5, 0x01, 0x0a, 0x0a, 0x41, + 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, + 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x12, 0x0a, + 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x79, 0x70, + 0x65, 0x12, 0x18, 0x0a, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x18, 0x03, 0x20, 0x01, + 0x28, 0x08, 0x52, 0x07, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x64, 0x12, 0x19, 0x0a, 0x08, 0x62, + 0x61, 0x73, 0x65, 0x5f, 0x75, 0x72, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, + 0x61, 0x73, 0x65, 0x55, 0x72, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x18, 0x05, + 0x20, 0x03, 0x28, 0x09, 0x52, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x12, 0x36, 0x0a, 0x07, 0x62, 0x65, + 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, + 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x52, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, + 0x63, 0x6b, 0x22, 0xf6, 0x01, 0x0a, 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x12, 0x16, 0x0a, 0x06, + 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, + 0x67, 0x69, 0x6f, 0x6e, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, + 0x65, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, + 0x4b, 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x11, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, 0x65, + 0x79, 0x5f, 0x73, 0x65, 0x63, 0x72, 0x65, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0f, + 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x4b, 0x65, 0x79, 0x53, 0x65, 0x63, 0x72, 0x65, 0x74, 0x12, + 0x14, 0x0a, 0x05, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, + 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x28, 0x0a, 0x10, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x5f, 0x66, + 0x61, 0x73, 0x74, 0x5f, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0e, 0x73, 0x6d, 0x61, 0x6c, 0x6c, 0x46, 0x61, 0x73, 0x74, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x12, + 0x19, 0x0a, 0x08, 0x72, 0x6f, 0x6c, 0x65, 0x5f, 0x61, 0x72, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x07, 0x72, 0x6f, 0x6c, 0x65, 0x41, 0x72, 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x65, 0x78, + 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x5f, 0x69, 0x64, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0a, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x49, 0x64, 0x32, 0xa9, 0x04, 0x0a, 0x08, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, + 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x12, 0x68, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, + 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, + 0x10, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, + 0x65, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x12, 0x56, 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, + 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, - 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, 0x63, - 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, 0x70, + 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, + 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, + 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, - 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, - 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, 0x52, - 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, - 0x74, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, 0x75, - 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, - 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, 0x6f, - 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, 0x65, + 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, + 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, + 0x68, 0x74, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, + 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, + 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, + 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, + 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, + 0x67, 0x73, 0x12, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, + 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, - 0x73, 0x12, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, - 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, - 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, - 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, 0x4d, - 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, - 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, - 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, - 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, - 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, - 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, 0x70, - 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xaa, 0x01, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, - 0x7a, 0x65, 0x72, 0x12, 0x47, 0x0a, 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, - 0x7a, 0x65, 0x64, 0x12, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, - 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, - 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, - 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, - 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, - 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, - 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x32, 0xbc, 0x01, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x43, 0x6f, - 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, 0x0a, 0x0e, 0x47, 0x65, - 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1c, 0x2e, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, - 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, - 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x55, 0x0a, 0x10, 0x57, 0x61, 0x74, - 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1e, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, - 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, - 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x30, 0x01, - 0x42, 0x32, 0x5a, 0x30, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, - 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, 0x6f, - 0x64, 0x65, 0x72, 0x64, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, + 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, + 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, + 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, + 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, + 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xaa, 0x01, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, + 0x69, 0x7a, 0x65, 0x72, 0x12, 0x47, 0x0a, 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, + 0x69, 0x7a, 0x65, 0x64, 0x12, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, + 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, + 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, + 0x10, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, + 0x64, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, + 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, + 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, + 0x73, 0x65, 0x32, 0xbc, 0x01, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x43, + 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, 0x0a, 0x0e, 0x47, + 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1c, 0x2e, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, + 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x55, 0x0a, 0x10, 0x57, 0x61, + 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1e, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, + 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, + 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x30, + 0x01, 0x42, 0x32, 0x5a, 0x30, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, + 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, + 0x6f, 0x64, 0x65, 0x72, 0x64, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( @@ -2224,43 +2211,42 @@ var file_coderd_aibridged_proto_aibridged_proto_depIdxs = []int32{ 14, // 12: proto.GetMCPServerConfigsResponse.external_auth_mcp_configs:type_name -> proto.MCPServerConfig 32, // 13: proto.GetMCPServerAccessTokensBatchResponse.access_tokens:type_name -> proto.GetMCPServerAccessTokensBatchResponse.AccessTokensEntry 33, // 14: proto.GetMCPServerAccessTokensBatchResponse.errors:type_name -> proto.GetMCPServerAccessTokensBatchResponse.ErrorsEntry - 34, // 15: proto.IsBudgetExceededRequest.period_start:type_name -> google.protobuf.Timestamp - 25, // 16: proto.GetAIProvidersResponse.providers:type_name -> proto.AIProvider - 26, // 17: proto.AIProvider.bedrock:type_name -> proto.AIProviderKindBedrock - 35, // 18: proto.RecordInterceptionRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 35, // 19: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 35, // 20: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 35, // 21: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 35, // 22: proto.RecordModelThoughtRequest.MetadataEntry.value:type_name -> google.protobuf.Any - 0, // 23: proto.Recorder.RecordInterception:input_type -> proto.RecordInterceptionRequest - 2, // 24: proto.Recorder.RecordInterceptionEnded:input_type -> proto.RecordInterceptionEndedRequest - 4, // 25: proto.Recorder.RecordTokenUsage:input_type -> proto.RecordTokenUsageRequest - 6, // 26: proto.Recorder.RecordPromptUsage:input_type -> proto.RecordPromptUsageRequest - 8, // 27: proto.Recorder.RecordToolUsage:input_type -> proto.RecordToolUsageRequest - 10, // 28: proto.Recorder.RecordModelThought:input_type -> proto.RecordModelThoughtRequest - 12, // 29: proto.MCPConfigurator.GetMCPServerConfigs:input_type -> proto.GetMCPServerConfigsRequest - 15, // 30: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:input_type -> proto.GetMCPServerAccessTokensBatchRequest - 17, // 31: proto.Authorizer.IsAuthorized:input_type -> proto.IsAuthorizedRequest - 19, // 32: proto.Authorizer.IsBudgetExceeded:input_type -> proto.IsBudgetExceededRequest - 21, // 33: proto.ProviderConfigurator.GetAIProviders:input_type -> proto.GetAIProvidersRequest - 23, // 34: proto.ProviderConfigurator.WatchAIProviders:input_type -> proto.WatchAIProvidersRequest - 1, // 35: proto.Recorder.RecordInterception:output_type -> proto.RecordInterceptionResponse - 3, // 36: proto.Recorder.RecordInterceptionEnded:output_type -> proto.RecordInterceptionEndedResponse - 5, // 37: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse - 7, // 38: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse - 9, // 39: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse - 11, // 40: proto.Recorder.RecordModelThought:output_type -> proto.RecordModelThoughtResponse - 13, // 41: proto.MCPConfigurator.GetMCPServerConfigs:output_type -> proto.GetMCPServerConfigsResponse - 16, // 42: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:output_type -> proto.GetMCPServerAccessTokensBatchResponse - 18, // 43: proto.Authorizer.IsAuthorized:output_type -> proto.IsAuthorizedResponse - 20, // 44: proto.Authorizer.IsBudgetExceeded:output_type -> proto.IsBudgetExceededResponse - 22, // 45: proto.ProviderConfigurator.GetAIProviders:output_type -> proto.GetAIProvidersResponse - 24, // 46: proto.ProviderConfigurator.WatchAIProviders:output_type -> proto.WatchAIProvidersResponse - 35, // [35:47] is the sub-list for method output_type - 23, // [23:35] is the sub-list for method input_type - 23, // [23:23] is the sub-list for extension type_name - 23, // [23:23] is the sub-list for extension extendee - 0, // [0:23] is the sub-list for field type_name + 25, // 15: proto.GetAIProvidersResponse.providers:type_name -> proto.AIProvider + 26, // 16: proto.AIProvider.bedrock:type_name -> proto.AIProviderKindBedrock + 35, // 17: proto.RecordInterceptionRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 35, // 18: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 35, // 19: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 35, // 20: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 35, // 21: proto.RecordModelThoughtRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 0, // 22: proto.Recorder.RecordInterception:input_type -> proto.RecordInterceptionRequest + 2, // 23: proto.Recorder.RecordInterceptionEnded:input_type -> proto.RecordInterceptionEndedRequest + 4, // 24: proto.Recorder.RecordTokenUsage:input_type -> proto.RecordTokenUsageRequest + 6, // 25: proto.Recorder.RecordPromptUsage:input_type -> proto.RecordPromptUsageRequest + 8, // 26: proto.Recorder.RecordToolUsage:input_type -> proto.RecordToolUsageRequest + 10, // 27: proto.Recorder.RecordModelThought:input_type -> proto.RecordModelThoughtRequest + 12, // 28: proto.MCPConfigurator.GetMCPServerConfigs:input_type -> proto.GetMCPServerConfigsRequest + 15, // 29: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:input_type -> proto.GetMCPServerAccessTokensBatchRequest + 17, // 30: proto.Authorizer.IsAuthorized:input_type -> proto.IsAuthorizedRequest + 19, // 31: proto.Authorizer.IsBudgetExceeded:input_type -> proto.IsBudgetExceededRequest + 21, // 32: proto.ProviderConfigurator.GetAIProviders:input_type -> proto.GetAIProvidersRequest + 23, // 33: proto.ProviderConfigurator.WatchAIProviders:input_type -> proto.WatchAIProvidersRequest + 1, // 34: proto.Recorder.RecordInterception:output_type -> proto.RecordInterceptionResponse + 3, // 35: proto.Recorder.RecordInterceptionEnded:output_type -> proto.RecordInterceptionEndedResponse + 5, // 36: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse + 7, // 37: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse + 9, // 38: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse + 11, // 39: proto.Recorder.RecordModelThought:output_type -> proto.RecordModelThoughtResponse + 13, // 40: proto.MCPConfigurator.GetMCPServerConfigs:output_type -> proto.GetMCPServerConfigsResponse + 16, // 41: proto.MCPConfigurator.GetMCPServerAccessTokensBatch:output_type -> proto.GetMCPServerAccessTokensBatchResponse + 18, // 42: proto.Authorizer.IsAuthorized:output_type -> proto.IsAuthorizedResponse + 20, // 43: proto.Authorizer.IsBudgetExceeded:output_type -> proto.IsBudgetExceededResponse + 22, // 44: proto.ProviderConfigurator.GetAIProviders:output_type -> proto.GetAIProvidersResponse + 24, // 45: proto.ProviderConfigurator.WatchAIProviders:output_type -> proto.WatchAIProvidersResponse + 34, // [34:46] is the sub-list for method output_type + 22, // [22:34] is the sub-list for method input_type + 22, // [22:22] is the sub-list for extension type_name + 22, // [22:22] is the sub-list for extension extendee + 0, // [0:22] is the sub-list for field type_name } func init() { file_coderd_aibridged_proto_aibridged_proto_init() } diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index 41795abdb9a..1bd3b3f6ffd 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -35,7 +35,7 @@ service Authorizer { // TODO: add authorization; currently only key validation takes place. rpc IsAuthorized(IsAuthorizedRequest) returns (IsAuthorizedResponse); // IsBudgetExceeded reports whether the user's AI spend has reached their - // effective limit over [period_start, now]. + // effective limit for the current deployment-configured budget period. rpc IsBudgetExceeded(IsBudgetExceededRequest) returns (IsBudgetExceededResponse); } @@ -189,8 +189,6 @@ message IsAuthorizedResponse { message IsBudgetExceededRequest { string user_id = 1; // UUID - // The spend aggregation window is [period_start, now]. - google.protobuf.Timestamp period_start = 2; } message IsBudgetExceededResponse { diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index cdfe739d027..af2d790e7ec 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -33,6 +33,7 @@ import ( coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" + "github.com/coder/quartz" ) var ( @@ -113,11 +114,15 @@ type Server struct { // budgetPolicy selects the effective group when a user belongs to multiple // budgeted groups, used for cost attribution on token usage records. budgetPolicy codersdk.AIBudgetPolicy + // budgetPeriod is the deployment-configured budgeting period used to + // derive the window over which user AI spend is aggregated. + budgetPeriod codersdk.AIBudgetPeriod + clock quartz.Clock } func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logger slog.Logger, accessURL string, bridgeCfg codersdk.AIBridgeConfig, externalAuthConfigs []*externalauth.Config, experiments codersdk.Experiments, - aiSeatTracker aiseats.SeatTracker, + aiSeatTracker aiseats.SeatTracker, clock quartz.Clock, ) (*Server, error) { eac := make(map[string]*externalauth.Config, len(externalAuthConfigs)) @@ -138,6 +143,8 @@ func NewServer(lifecycleCtx context.Context, store store, ps pubsub.Pubsub, logg structuredLogging: bridgeCfg.StructuredLogging.Value(), aiSeatTracker: aiSeatTracker, budgetPolicy: codersdk.NewAIBudgetPolicyFromString(bridgeCfg.BudgetPolicy), + budgetPeriod: codersdk.NewAIBudgetPeriodFromString(bridgeCfg.BudgetPeriod), + clock: clock, } if bridgeCfg.InjectCoderMCPTools { @@ -747,7 +754,8 @@ 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]. +// 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) { //nolint:gocritic // AIBridged has specific authz rules. ctx = dbauthz.AsAIBridged(ctx) @@ -756,14 +764,13 @@ func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceede if err != nil { return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err) } - // An unset PeriodStart deserializes to time.Unix(0, 0), which would - // incorrectly aggregate the user's lifetime spend against a period budget. - if in.PeriodStart == nil { - return nil, xerrors.New("period_start is required") + + periodWindow, err := budget.CurrentPeriod(s.clock.Now(), s.budgetPeriod) + if err != nil { + return nil, xerrors.Errorf("compute AI budget period: %w", err) } - periodStart := in.GetPeriodStart().AsTime() - userBudget, err := s.checkUserAIBudget(ctx, userID, periodStart) + userBudget, err := s.checkUserAIBudget(ctx, userID, periodWindow.Start) if err != nil { return nil, err } diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index b4eec74b663..e5b71803eb9 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -29,7 +29,6 @@ import ( "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogjson" "cdr.dev/slog/v3/sloggers/slogtest" - "github.com/coder/coder/v2/coderd/aibridge/budget" "github.com/coder/coder/v2/coderd/aibridged" "github.com/coder/coder/v2/coderd/aibridged/proto" "github.com/coder/coder/v2/coderd/aibridgedserver" @@ -51,6 +50,7 @@ import ( "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" "github.com/coder/serpent" ) @@ -209,7 +209,7 @@ func TestAuthorization(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) require.NotNil(t, srv) @@ -371,7 +371,7 @@ func TestAuthorization_Delegated(t *testing.T) { tc.mocksFn(db, apiKey, user) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) require.NotNil(t, srv) @@ -402,7 +402,6 @@ func TestIsBudgetExceeded(t *testing.T) { cases := []struct { name string userIDStr string - omitPeriodStart bool setupMocks func(db *dbmock.MockStore, userID uuid.UUID) *proto.IsBudgetExceededResponse wantErrContains string }{ @@ -412,12 +411,6 @@ func TestIsBudgetExceeded(t *testing.T) { userIDStr: "not-a-uuid", wantErrContains: "invalid user_id", }, - { - // Missing period_start is rejected before any store call. - name: "missing period_start", - omitPeriodStart: true, - wantErrContains: "period_start is required", - }, { // No override and no group budget resolves: pass-through. name: "no budget configured returns not exceeded", @@ -567,15 +560,10 @@ func TestIsBudgetExceeded(t *testing.T) { wantResp = tc.setupMocks(db, userID) } - srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) req := &proto.IsBudgetExceededRequest{UserId: userIDStr} - if !tc.omitPeriodStart { - window, err := budget.CurrentPeriod(dbtime.Now(), codersdk.AIBudgetPeriodMonth) - require.NoError(t, err) - req.PeriodStart = timestamppb.New(window.Start) - } resp, err := srv.IsBudgetExceeded(t.Context(), req) if tc.wantErrContains != "" { require.Error(t, err) @@ -599,7 +587,7 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { const groupLimitMicros = 1_000_000 // setup provisions a user in an organization with a single budgeted group. - setup := func(t *testing.T) (context.Context, database.Store, *aibridgedserver.Server, database.User, database.Group) { + setup := func(t *testing.T, clock quartz.Clock) (context.Context, database.Store, *aibridgedserver.Server, database.User, database.Group) { t.Helper() ctx := testutil.Context(t, testutil.WaitLong) @@ -620,7 +608,7 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { }) require.NoError(t, err) - srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, clock) require.NoError(t, err) return ctx, rawDB, srv, user, group @@ -628,59 +616,68 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { t.Run("period boundary excludes prior period spend", func(t *testing.T) { t.Parallel() - ctx, rawDB, srv, user, group := setup(t) - prevMonth := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) - newMonth := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC) + // Use fixed dates to keep the test deterministic. + clock := quartz.NewMock(t) + ctx, rawDB, srv, user, group := setup(t, clock) - // User spend on 2026-01-15. + prevMonth := time.Date(2026, time.January, 15, 0, 0, 0, 0, time.UTC) + nextMonth := time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC) + + // Set now to 2026-01-15. + clock.Set(prevMonth) + + // User spend on 2026-01-15 exceeds the group limit. _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, - Day: prevMonth.AddDate(0, 0, 14), + Day: prevMonth, CostMicros: 1_500_000, }) require.NoError(t, err) - // Query with period_start 2026-01-01: includes the 2026-01-15 spend, user exceeded. + // Current period is January: includes the 2026-01-15 spend, user exceeded. prevMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(prevMonth), + UserId: user.ID.String(), }) require.NoError(t, err) require.True(t, prevMonthResp.GetExceeded()) require.Equal(t, int64(groupLimitMicros), prevMonthResp.GetSpendLimitMicros()) - // Query with period_start 2026-02-01: excludes the 2026-01-15 spend, user not exceeded. - newMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(newMonth), + // Advance clock to 2026-02-05: excludes the 2026-01-15 spend, user not exceeded. + clock.Set(nextMonth) + nextMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ + UserId: user.ID.String(), }) require.NoError(t, err) - require.False(t, newMonthResp.GetExceeded()) - require.Equal(t, int64(groupLimitMicros), newMonthResp.GetSpendLimitMicros()) + require.False(t, nextMonthResp.GetExceeded()) + require.Equal(t, int64(groupLimitMicros), nextMonthResp.GetSpendLimitMicros()) }) t.Run("new user override unblocks user", func(t *testing.T) { t.Parallel() - ctx, rawDB, srv, user, group := setup(t) // Use fixed dates to keep the test deterministic. - periodStart := time.Date(2026, time.March, 1, 0, 0, 0, 0, time.UTC) + clock := quartz.NewMock(t) + ctx, rawDB, srv, user, group := setup(t, clock) + + now := time.Date(2026, time.March, 15, 0, 0, 0, 0, time.UTC) + + // Set now to 2026-03-15. + clock.Set(now) - // User spend on 2026-03-15. + // User spend on 2026-03-15 exceeds the group limit. _, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{ UserID: user.ID, EffectiveGroupID: group.ID, - Day: periodStart.AddDate(0, 0, 14), + Day: now, CostMicros: 1_500_000, }) require.NoError(t, err) // User's spend exceeds the group limit. beforeResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(periodStart), + UserId: user.ID.String(), }) require.NoError(t, err) require.True(t, beforeResp.GetExceeded()) @@ -697,8 +694,7 @@ func TestIsBudgetExceeded_Enforcement(t *testing.T) { require.NoError(t, err) afterResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{ - UserId: user.ID.String(), - PeriodStart: timestamppb.New(periodStart), + UserId: user.ID.String(), }) require.NoError(t, err) require.False(t, afterResp.GetExceeded()) @@ -777,7 +773,7 @@ func TestGetMCPServerConfigs(t *testing.T) { accessURL := "https://my-cool-deployment.com" srv, err := aibridgedserver.NewServer(t.Context(), db, nil, logger, accessURL, codersdk.AIBridgeConfig{ InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection), - }, tc.externalAuthConfigs, tc.experiments, agplaiseats.Noop{}) + }, tc.externalAuthConfigs, tc.experiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) require.NotNil(t, srv) @@ -827,7 +823,7 @@ func TestGetMCPServerAccessTokensBatch(t *testing.T) { { ID: "3", }, - }, requiredExperiments, agplaiseats.Noop{}) + }, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) require.NotNil(t, srv) @@ -2104,7 +2100,7 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC) // The server runs every store call as subjectAibridged via the authzDB. - srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, authzDB, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) _, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ @@ -2479,7 +2475,7 @@ func testRecordMethod[Req any, Resp any]( } ctx := testutil.Context(t, testutil.WaitLong) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) resp, err := callMethod(srv, ctx, tc.request) @@ -2801,7 +2797,7 @@ func TestStructuredLogging(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{ StructuredLogging: serpent.Bool(tc.structuredLogging), - }, nil, requiredExperiments, agplaiseats.Noop{}) + }, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) err = tc.recordFn(srv, ctx, interceptionID) @@ -2843,7 +2839,7 @@ func TestInferredThreadsByToolCalls(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) aID := uuid.New() @@ -2939,7 +2935,7 @@ func TestRecordToolUsageProviderItemID(t *testing.T) { user := dbgen.User(t, db, database.User{}) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) intcID := uuid.New() @@ -3069,7 +3065,7 @@ func TestGetAIProviders(t *testing.T) { Settings: sql.NullString{String: "{not valid json", Valid: true}, }) - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) resp, err := srv.GetAIProviders(ctx, &proto.GetAIProvidersRequest{}) @@ -3132,7 +3128,7 @@ func TestGetAIProvidersBlocksOnSeedLock(t *testing.T) { BaseUrl: "https://api.openai.com/", }, "sk-openai") - srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, nil, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) // Simulate an in-flight env seed holding the advisory lock until released. @@ -3213,7 +3209,7 @@ func TestWatchAIProviders(t *testing.T) { // In-memory pubsub delivers Publish synchronously for deterministic signals. ps := pubsub.NewInMemory() - srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) @@ -3251,7 +3247,7 @@ func TestWatchAIProvidersSignalsOnDeliveryError(t *testing.T) { logger := slogtest.Make(t, nil) ps := &captureListenerPubsub{listenerC: make(chan pubsub.ListenerWithErr, 1)} - srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(ctx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) @@ -3293,7 +3289,7 @@ func TestWatchAIProvidersStopsOnLifecycleCancel(t *testing.T) { // canceled while the stream stays open. lifecycleCtx, lifecycleCancel := context.WithCancel(ctx) defer lifecycleCancel() - srv, err := aibridgedserver.NewServer(lifecycleCtx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}) + srv, err := aibridgedserver.NewServer(lifecycleCtx, db, ps, logger, "/", codersdk.AIBridgeConfig{}, nil, nil, agplaiseats.Noop{}, quartz.NewReal()) require.NoError(t, err) streamCtx, streamCancel := context.WithCancel(ctx) diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 2edd295dbc6..8658ef35039 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -612,6 +612,15 @@ var AIBudgetPeriods = []string{ string(AIBudgetPeriodMonth), } +// NewAIBudgetPeriodFromString converts s to an AIBudgetPeriod, falling back to +// AIBudgetPeriodMonth when s is empty or not a recognized period. +func NewAIBudgetPeriodFromString(s string) AIBudgetPeriod { + if slices.Contains(AIBudgetPeriods, s) { + return AIBudgetPeriod(s) + } + return AIBudgetPeriodMonth +} + // DeploymentValues is the central configuration values the coder server. type DeploymentValues struct { Verbose serpent.Bool `json:"verbose,omitempty"` diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 8736088a803..ac6b3d80fe1 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -890,7 +890,8 @@ func (api *API) userAISpendStatus(rw http.ResponseWriter, r *http.Request) { user := httpmw.UserParam(r) logger := api.Logger.With(slog.F("user_id", user.ID)) - periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), codersdk.AIBudgetPeriodMonth) + period := codersdk.NewAIBudgetPeriodFromString(api.DeploymentValues.AI.BridgeConfig.BudgetPeriod) + periodWindow, err := budget.CurrentPeriod(api.Clock.Now(), period) if err != nil { logger.Error(ctx, "failed to compute AI budget period", slog.Error(err)) httpapi.InternalServerError(rw, err) diff --git a/enterprise/coderd/aibridgeserve.go b/enterprise/coderd/aibridgeserve.go index 08b86f4cd62..8e90ae6e71a 100644 --- a/enterprise/coderd/aibridgeserve.go +++ b/enterprise/coderd/aibridgeserve.go @@ -144,6 +144,7 @@ func (api *API) aiGatewayServe(rw http.ResponseWriter, r *http.Request) { api.ExternalAuthConfigs, api.AGPL.Experiments, api.AGPL.AISeatTracker, + api.AGPL.Clock, ) if err != nil { if !xerrors.Is(err, context.Canceled) {