feat: public endpoint

toeverything · Apr 10, 2024 · 2199bad · 2199bad
1 parent 35a4768
commit 2199bad
Show file tree

Hide file tree

Showing 15 changed files with 253 additions and 48 deletions.
diff --git a/.github/actions/deploy/deploy.mjs b/.github/actions/deploy/deploy.mjs
@@ -118,6 +118,7 @@ const createHelmCommand = ({ isDryRun }) => {
     `--set        graphql.app.experimental.enableJwstCodec=${namespace === 'dev'}`,
     `--set        graphql.app.features.earlyAccessPreview=false`,
     `--set        graphql.app.features.syncClientVersionCheck=true`,
+    `--set        graphql.app.features.copilotAuthorization=true`,
     `--set        sync.replicaCount=${syncReplicaCount}`,
     `--set-string sync.image.tag="${imageTag}"`,
     ...serviceAnnotations,

diff --git a/.github/helm/affine/charts/graphql/templates/deployment.yaml b/.github/helm/affine/charts/graphql/templates/deployment.yaml
@@ -83,6 +83,8 @@ spec:
             value: "{{ .Values.app.features.earlyAccessPreview }}"
           - name: FEATURES_SYNC_CLIENT_VERSION_CHECK
             value: "{{ .Values.app.features.syncClientVersionCheck }}"
+          - name: FEATURES_COPILOT_SKIP_AUTHORIZATION
+            value: "{{ .Values.app.features.copilotAuthorization }}"
           - name: MAILER_HOST
             valueFrom:
               secretKeyRef:

diff --git a/.github/helm/affine/charts/graphql/values.yaml b/.github/helm/affine/charts/graphql/values.yaml
@@ -62,6 +62,7 @@ app:
   features:
     earlyAccessPreview: false
     syncClientVersionCheck: false
+    copilotAuthorization: false
 
 serviceAccount:
   create: true

diff --git a/packages/backend/server/src/config/affine.env.ts b/packages/backend/server/src/config/affine.env.ts
@@ -38,5 +38,9 @@ AFFiNE.ENV_MAP = {
     'featureFlags.syncClientVersionCheck',
     'boolean',
   ],
+  FEATURES_COPILOT_SKIP_AUTHORIZATION: [
+    'featureFlags.copilotAuthorization',
+    'boolean',
+  ],
   TELEMETRY_ENABLE: ['telemetry.enabled', 'boolean'],
 };
diff --git a/packages/backend/server/src/fundamentals/config/def.ts b/packages/backend/server/src/fundamentals/config/def.ts
@@ -185,6 +185,7 @@ export interface AFFiNEConfig {
   featureFlags: {
     earlyAccessPreview: boolean;
     syncClientVersionCheck: boolean;
+    copilotAuthorization: boolean;
   };
 
   /**

diff --git a/packages/backend/server/src/fundamentals/config/default.ts b/packages/backend/server/src/fundamentals/config/default.ts
@@ -119,6 +119,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
     featureFlags: {
       earlyAccessPreview: false,
       syncClientVersionCheck: false,
+      copilotAuthorization: false,
     },
     https: false,
     host: 'localhost',

diff --git a/packages/backend/server/src/fundamentals/graphql/index.ts b/packages/backend/server/src/fundamentals/graphql/index.ts
@@ -23,6 +23,15 @@ export type GraphqlContext = {
     GraphQLModule.forRootAsync<ApolloDriverConfig>({
       driver: ApolloDriver,
       useFactory: (config: Config) => {
+        const copilotAuthorization = config.featureFlags.copilotAuthorization;
+        const cors = {
+          cors: {
+            origin: [
+              'https://try-blocksuite.vercel.app/',
+              'http://localhost:5173/',
+            ],
+          },
+        };
         return {
           ...config.graphql,
           path: `${config.path}/graphql`,
@@ -78,6 +87,7 @@ export type GraphqlContext = {
 
             return formattedError;
           },
+          ...(copilotAuthorization ? cors : {}),
         };
       },
       inject: [Config],

diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts
@@ -82,7 +82,7 @@ export class CopilotController {
   @Public()
   @Get('/chat/:sessionId')
   async chat(
-    @CurrentUser() user: CurrentUser,
+    @CurrentUser() user: CurrentUser | undefined,
     @Req() req: Request,
     @Param('sessionId') sessionId: string,
     @Query('message') message: string | undefined,
@@ -110,7 +110,7 @@ export class CopilotController {
         session.model,
         {
           signal: req.signal,
-          user: user.id,
+          user: user?.id,
         }
       );
 
@@ -132,7 +132,7 @@ export class CopilotController {
   @Public()
   @Sse('/chat/:sessionId/stream')
   async chatStream(
-    @CurrentUser() user: CurrentUser,
+    @CurrentUser() user: CurrentUser | undefined,
     @Req() req: Request,
     @Param('sessionId') sessionId: string,
     @Query('message') message: string | undefined,
@@ -157,7 +157,7 @@ export class CopilotController {
     return from(
       provider.generateTextStream(session.finish(params), session.model, {
         signal: req.signal,
-        user: user.id,
+        user: user?.id,
       })
     ).pipe(
       connect(shared$ =>

diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts
@@ -1,4 +1,4 @@
-import { Logger } from '@nestjs/common';
+import { ForbiddenException, Logger } from '@nestjs/common';
 import {
   Args,
   Field,
@@ -7,13 +7,14 @@ import {
   Mutation,
   ObjectType,
   Parent,
+  Query,
   registerEnumType,
   ResolveField,
   Resolver,
 } from '@nestjs/graphql';
 import { SafeIntResolver } from 'graphql-scalars';
 
-import { CurrentUser } from '../../core/auth';
+import { CurrentUser, Public } from '../../core/auth';
 import { QuotaService } from '../../core/quota';
 import { UserType } from '../../core/user';
 import { PermissionService } from '../../core/workspaces/permission';
@@ -158,8 +159,11 @@ export class CopilotResolver {
   })
   async getQuota(
     @Parent() _copilot: CopilotType,
-    @CurrentUser() user: CurrentUser
+    @CurrentUser() user: CurrentUser | undefined
   ) {
+    // TODO(@darkskygit): remove this after the feature is stable
+    if (!user) return { used: 0 };
+
     const quota = await this.quota.getUserQuota(user.id);
     const limit = quota.feature.copilotActionLimit;
 
@@ -182,11 +186,11 @@ export class CopilotResolver {
   })
   async chats(
     @Parent() copilot: CopilotType,
-    @CurrentUser() user: CurrentUser
+    @CurrentUser() user?: CurrentUser
   ) {
     if (!copilot.workspaceId) return [];
-    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
-    return await this.chatSession.listSessions(user.id, copilot.workspaceId);
+    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user?.id);
+    return await this.chatSession.listSessions(user?.id);
   }
 
   @ResolveField(() => [String], {
@@ -195,19 +199,19 @@ export class CopilotResolver {
   })
   async actions(
     @Parent() copilot: CopilotType,
-    @CurrentUser() user: CurrentUser
+    @CurrentUser() user?: CurrentUser
   ) {
     if (!copilot.workspaceId) return [];
-    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
-    return await this.chatSession.listSessions(user.id, copilot.workspaceId, {
+    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user?.id);
+    return await this.chatSession.listSessions(user?.id, copilot.workspaceId, {
       action: true,
     });
   }
 
   @ResolveField(() => [CopilotHistoriesType], {})
   async histories(
     @Parent() copilot: CopilotType,
-    @CurrentUser() user: CurrentUser,
+    @CurrentUser() user?: CurrentUser,
     @Args('docId', { nullable: true }) docId?: string,
     @Args({
       name: 'options',
@@ -217,70 +221,88 @@ export class CopilotResolver {
     options?: QueryChatHistoriesInput
   ) {
     const workspaceId = copilot.workspaceId;
-    if (!workspaceId) {
-      return [];
-    } else if (docId) {
-      await this.permissions.checkCloudPagePermission(
-        workspaceId,
-        docId,
-        user.id
-      );
-    } else {
-      await this.permissions.checkCloudWorkspace(workspaceId, user.id);
+    // todo(@darkskygit): remove this after the feature is stable
+    const publishable = AFFiNE.featureFlags.copilotAuthorization;
+    if (user) {
+      if (!workspaceId) {
+        return [];
+      }
+      if (docId) {
+        await this.permissions.checkCloudPagePermission(
+          workspaceId,
+          docId,
+          user.id
+        );
+      } else {
+        await this.permissions.checkCloudWorkspace(workspaceId, user.id);
+      }
+    } else if (!publishable) {
+      return new ForbiddenException('Login required');
     }
 
     return await this.chatSession.listHistories(
-      user.id,
+      user?.id,
       workspaceId,
       docId,
       options
     );
   }
 
+  @Public()
   @Mutation(() => String, {
     description: 'Create a chat session',
   })
   async createCopilotSession(
-    @CurrentUser() user: CurrentUser,
+    @CurrentUser() user: CurrentUser | undefined,
     @Args({ name: 'options', type: () => CreateChatSessionInput })
     options: CreateChatSessionInput
   ) {
-    await this.permissions.checkCloudPagePermission(
-      options.workspaceId,
-      options.docId,
-      user.id
-    );
-    const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
+    // todo(@darkskygit): remove this after the feature is stable
+    const publishable = AFFiNE.featureFlags.copilotAuthorization;
+    if (!user && !publishable) {
+      return new ForbiddenException('Login required');
+    }
+
+    const lockFlag = `${COPILOT_LOCKER}:session:${user?.id}:${options.workspaceId}`;
     await using lock = await this.mutex.lock(lockFlag);
     if (!lock) {
       return new TooManyRequestsException('Server is busy');
     }
-
-    const { limit, used } = await this.getQuota(
-      { workspaceId: undefined },
-      user
-    );
-    if (limit && Number.isFinite(limit) && used >= limit) {
-      return new PaymentRequiredException(
-        `You have reached the limit of actions in this workspace, please upgrade your plan.`
+    if (options.action && user) {
+      const { limit, used } = await this.getQuota(
+        { workspaceId: undefined },
+        user
       );
+      if (limit && Number.isFinite(limit) && used >= limit) {
+        return new PaymentRequiredException(
+          `You have reached the limit of actions in this workspace, please upgrade your plan.`
+        );
+      }
     }
 
     const session = await this.chatSession.create({
       ...options,
-      userId: user.id,
+      // todo: force user to be logged in
+      userId: user?.id ?? '',
     });
     return session;
   }
 
+  @Public()
   @Mutation(() => String, {
     description: 'Create a chat message',
   })
   async createCopilotMessage(
-    @CurrentUser() user: CurrentUser,
+    @CurrentUser() user: CurrentUser | undefined,
     @Args({ name: 'options', type: () => CreateChatMessageInput })
     options: CreateChatMessageInput
   ) {
+    // todo(@darkskygit): remove this after the feature is stable
+    const publishable = AFFiNE.featureFlags.copilotAuthorization;
+    if (!user && !publishable) {
+      return new ForbiddenException('Login required');
+    }
+
     const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`;
     await using lock = await this.mutex.lock(lockFlag);
     if (!lock) {
@@ -314,4 +336,11 @@ export class UserCopilotResolver {
     }
     return { workspaceId };
   }
+
+  @Public()
+  @Query(() => CopilotType)
+  async copilotAnonymous(@Args('workspaceId') workspaceId: string) {
+    if (!AFFiNE.featureFlags.copilotAuthorization) return;
+    return { workspaceId };
+  }
 }