feat(server): support registering ai early access users (#6565)

packages/backend/server/src/core/features/feature.ts

@@ -61,7 +61,18 @@ export class UnlimitedCopilotFeatureConfig extends FeatureConfig {
     super(data);
 
     if (this.config.feature !== FeatureType.UnlimitedCopilot) {
-      throw new Error('Invalid feature config: type is not UnlimitedWorkspace');
+      throw new Error('Invalid feature config: type is not AIEarlyAccess');
+    }
+  }
+}
+export class AIEarlyAccessFeatureConfig extends FeatureConfig {
+  override config!: Feature & { feature: FeatureType.AIEarlyAccess };
+
+  constructor(data: any) {
+    super(data);
+
+    if (this.config.feature !== FeatureType.AIEarlyAccess) {
+      throw new Error('Invalid feature config: type is not AIEarlyAccess');
     }
   }
 }
@@ -69,6 +80,7 @@ export class UnlimitedCopilotFeatureConfig extends FeatureConfig {
 const FeatureConfigMap = {
   [FeatureType.Copilot]: CopilotFeatureConfig,
   [FeatureType.EarlyAccess]: EarlyAccessFeatureConfig,
+  [FeatureType.AIEarlyAccess]: AIEarlyAccessFeatureConfig,
   [FeatureType.UnlimitedWorkspace]: UnlimitedWorkspaceFeatureConfig,
   [FeatureType.UnlimitedCopilot]: UnlimitedCopilotFeatureConfig,
 };

packages/backend/server/src/core/features/index.ts

@@ -1,6 +1,6 @@
 import { Module } from '@nestjs/common';
 
-import { FeatureManagementService } from './management';
+import { EarlyAccessType, FeatureManagementService } from './management';
 import { FeatureService } from './service';
 
 /**
@@ -15,6 +15,11 @@ import { FeatureService } from './service';
 })
 export class FeatureModule {}
 
-export { type CommonFeature, commonFeatureSchema } from './types';
-export { FeatureKind, Features, FeatureType } from './types';
-export { FeatureManagementService, FeatureService };
+export {
+  type CommonFeature,
+  commonFeatureSchema,
+  FeatureKind,
+  Features,
+  FeatureType,
+} from './types';
+export { EarlyAccessType, FeatureManagementService, FeatureService };

packages/backend/server/src/core/features/management.ts

@@ -7,6 +7,11 @@ import { FeatureType } from './types';
 
 const STAFF = ['@toeverything.info'];
 
+export enum EarlyAccessType {
+  App = 'app',
+  AI = 'ai',
+}
+
 @Injectable()
 export class FeatureManagementService {
   protected logger = new Logger(FeatureManagementService.name);
@@ -30,24 +35,43 @@ export class FeatureManagementService {
   }
 
   // ======== Early Access ========
-
-  async addEarlyAccess(userId: string) {
+  async addEarlyAccess(
+    userId: string,
+    type: EarlyAccessType = EarlyAccessType.App
+  ) {
     return this.feature.addUserFeature(
       userId,
-      FeatureType.EarlyAccess,
+      type === EarlyAccessType.App
+        ? FeatureType.EarlyAccess
+        : FeatureType.AIEarlyAccess,
       'Early access user'
     );
   }
 
-  async removeEarlyAccess(userId: string) {
-    return this.feature.removeUserFeature(userId, FeatureType.EarlyAccess);
+  async removeEarlyAccess(
+    userId: string,
+    type: EarlyAccessType = EarlyAccessType.App
+  ) {
+    return this.feature.removeUserFeature(
+      userId,
+      type === EarlyAccessType.App
+        ? FeatureType.EarlyAccess
+        : FeatureType.AIEarlyAccess
+    );
   }
 
-  async listEarlyAccess() {
-    return this.feature.listFeatureUsers(FeatureType.EarlyAccess);
+  async listEarlyAccess(type: EarlyAccessType = EarlyAccessType.App) {
+    return this.feature.listFeatureUsers(
+      type === EarlyAccessType.App
+        ? FeatureType.EarlyAccess
+        : FeatureType.AIEarlyAccess
+    );
   }
 
-  async isEarlyAccessUser(email: string) {
+  async isEarlyAccessUser(
+    email: string,
+    type: EarlyAccessType = EarlyAccessType.App
+  ) {
     const user = await this.prisma.user.findFirst({
       where: {
         email: {
@@ -56,9 +80,15 @@ export class FeatureManagementService {
         },
       },
     });
+
     if (user) {
       const canEarlyAccess = await this.feature
-        .hasUserFeature(user.id, FeatureType.EarlyAccess)
+        .hasUserFeature(
+          user.id,
+          type === EarlyAccessType.App
+            ? FeatureType.EarlyAccess
+            : FeatureType.AIEarlyAccess
+        )
         .catch(() => false);
 
       return canEarlyAccess;
@@ -67,9 +97,12 @@ export class FeatureManagementService {
   }
 
   /// check early access by email
-  async canEarlyAccess(email: string) {
+  async canEarlyAccess(
+    email: string,
+    type: EarlyAccessType = EarlyAccessType.App
+  ) {
     if (this.config.featureFlags.earlyAccessPreview && !this.isStaff(email)) {
-      return this.isEarlyAccessUser(email);
+      return this.isEarlyAccessUser(email, type);
     } else {
       return true;
     }

packages/backend/server/src/core/features/service.ts

@@ -63,13 +63,6 @@ export class FeatureService {
     expiredAt?: Date | string
   ) {
     return this.prisma.$transaction(async tx => {
-      const latestVersion = await tx.features
-        .aggregate({
-          where: { feature },
-          _max: { version: true },
-        })
-        .then(r => r._max.version || 1);
-
       const latestFlag = await tx.userFeatures.findFirst({
         where: {
           userId,
@@ -83,9 +76,21 @@ export class FeatureService {
           createdAt: 'desc',
         },
       });
+
       if (latestFlag) {
         return latestFlag.id;
       } else {
+        const latestVersion = await tx.features
+          .aggregate({
+            where: { feature },
+            _max: { version: true },
+          })
+          .then(r => r._max.version);
+
+        if (!latestVersion) {
+          throw new Error(`Feature ${feature} not found`);
+        }
+
         return tx.userFeatures
           .create({
             data: {

packages/backend/server/src/core/features/types/common.ts

@@ -3,6 +3,7 @@ import { registerEnumType } from '@nestjs/graphql';
 export enum FeatureType {
   // user feature
   EarlyAccess = 'early_access',
+  AIEarlyAccess = 'ai_early_access',
   UnlimitedCopilot = 'unlimited_copilot',
   // workspace feature
   Copilot = 'copilot',

packages/backend/server/src/core/features/types/early-access.ts

@@ -9,3 +9,8 @@ export const featureEarlyAccess = z.object({
     whitelist: z.string().array(),
   }),
 });
+
+export const featureAIEarlyAccess = z.object({
+  feature: z.literal(FeatureType.AIEarlyAccess),
+  configs: z.object({}),
+});

packages/backend/server/src/core/features/types/index.ts

@@ -2,7 +2,7 @@ import { z } from 'zod';
 
 import { FeatureType } from './common';
 import { featureCopilot } from './copilot';
-import { featureEarlyAccess } from './early-access';
+import { featureAIEarlyAccess, featureEarlyAccess } from './early-access';
 import { featureUnlimitedCopilot } from './unlimited-copilot';
 import { featureUnlimitedWorkspace } from './unlimited-workspace';
 
@@ -59,6 +59,12 @@ export const Features: Feature[] = [
     version: 1,
     configs: {},
   },
+  {
+    feature: FeatureType.AIEarlyAccess,
+    type: FeatureKind.Feature,
+    version: 1,
+    configs: {},
+  },
 ];
 
 /// ======== schema infer ========
@@ -71,6 +77,7 @@ export const FeatureSchema = commonFeatureSchema
     z.discriminatedUnion('feature', [
       featureCopilot,
       featureEarlyAccess,
+      featureAIEarlyAccess,
       featureUnlimitedWorkspace,
       featureUnlimitedCopilot,
     ])

packages/backend/server/src/core/user/management.ts

@@ -3,15 +3,27 @@ import {
   ForbiddenException,
   UseGuards,
 } from '@nestjs/common';
-import { Args, Context, Int, Mutation, Query, Resolver } from '@nestjs/graphql';
+import {
+  Args,
+  Context,
+  Int,
+  Mutation,
+  Query,
+  registerEnumType,
+  Resolver,
+} from '@nestjs/graphql';
 
 import { CloudThrottlerGuard, Throttle } from '../../fundamentals';
 import { CurrentUser } from '../auth/current-user';
 import { sessionUser } from '../auth/service';
-import { FeatureManagementService } from '../features';
+import { EarlyAccessType, FeatureManagementService } from '../features';
 import { UserService } from './service';
 import { UserType } from './types';
 
+registerEnumType(EarlyAccessType, {
+  name: 'EarlyAccessType',
+});
+
 /**
  * User resolver
  * All op rate limit: 10 req/m
@@ -33,19 +45,20 @@ export class UserManagementResolver {
   @Mutation(() => Int)
   async addToEarlyAccess(
     @CurrentUser() currentUser: CurrentUser,
-    @Args('email') email: string
+    @Args('email') email: string,
+    @Args({ name: 'type', type: () => EarlyAccessType }) type: EarlyAccessType
   ): Promise<number> {
     if (!this.feature.isStaff(currentUser.email)) {
       throw new ForbiddenException('You are not allowed to do this');
     }
     const user = await this.users.findUserByEmail(email);
     if (user) {
-      return this.feature.addEarlyAccess(user.id);
+      return this.feature.addEarlyAccess(user.id, type);
     } else {
       const user = await this.users.createAnonymousUser(email, {
         registered: false,
       });
-      return this.feature.addEarlyAccess(user.id);
+      return this.feature.addEarlyAccess(user.id, type);
     }
   }