refactor(core): get copilot sessions api (#10168)

Fix issue [BS-2575](https://linear.app/affine-design/issue/BS-2575).

### What Changed?
- Refactor `getCopilotSessions` api.
  - Add `docId` parameter.
  - Add `action` parameter.

packages/backend/server/src/plugins/copilot/resolver.ts

@@ -129,6 +129,12 @@ enum ChatHistoryOrder {
 
 registerEnumType(ChatHistoryOrder, { name: 'ChatHistoryOrder' });
 
+@InputType()
+class QueryChatSessionsInput {
+  @Field(() => Boolean, { nullable: true })
+  action: boolean | undefined;
+}
+
 @InputType()
 class QueryChatHistoriesInput implements Partial<ListHistoriesOptions> {
   @Field(() => Boolean, { nullable: true })
@@ -274,6 +280,9 @@ class CopilotPromptType {
 export class CopilotType {
   @Field(() => ID, { nullable: true })
   workspaceId!: string | undefined;
+
+  @Field(() => ID, { nullable: true })
+  docId!: string | undefined;
 }
 
 @Throttle()
@@ -296,31 +305,23 @@ export class CopilotResolver {
   }
 
   @ResolveField(() => [String], {
-    description: 'Get the session list of chats in the workspace',
+    description: 'Get the session list in the workspace',
     complexity: 2,
   })
-  async chats(
+  async sessionIds(
     @Parent() copilot: CopilotType,
-    @CurrentUser() user: CurrentUser
+    @CurrentUser() user: CurrentUser,
+    @Args('docId', { nullable: true }) docId?: string,
+    @Args('options', { nullable: true }) options?: QueryChatSessionsInput
   ) {
     if (!copilot.workspaceId) return [];
     await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
-    return await this.chatSession.listSessions(user.id, copilot.workspaceId);
-  }
-
-  @ResolveField(() => [String], {
-    description: 'Get the session list of actions in the workspace',
-    complexity: 2,
-  })
-  async actions(
-    @Parent() copilot: CopilotType,
-    @CurrentUser() user: CurrentUser
-  ) {
-    if (!copilot.workspaceId) return [];
-    await this.permissions.checkCloudWorkspace(copilot.workspaceId, user.id);
-    return await this.chatSession.listSessions(user.id, copilot.workspaceId, {
-      action: true,
-    });
+    return await this.chatSession.listSessionIds(
+      user.id,
+      copilot.workspaceId,
+      docId,
+      options
+    );
   }
 
   @ResolveField(() => [CopilotHistoriesType], {})

packages/backend/server/src/plugins/copilot/session.ts

@@ -391,17 +391,18 @@ export class ChatSessionService {
       .reduce((prev, cost) => prev + cost, 0);
   }
 
-  async listSessions(
+  async listSessionIds(
     userId: string,
     workspaceId: string,
-    options?: { docId?: string; action?: boolean }
+    docId?: string,
+    options?: { action?: boolean }
   ): Promise<string[]> {
     return await this.db.aiSession
       .findMany({
         where: {
           userId,
           workspaceId,
-          docId: workspaceId === options?.docId ? undefined : options?.docId,
+          docId,
           prompt: {
             action: options?.action ? { not: null } : null,
           },

packages/backend/server/src/schema.gql

@@ -37,18 +37,16 @@ enum ContextFileStatus {
 }
 
 type Copilot {
-  """Get the session list of actions in the workspace"""
-  actions: [String!]!
-
-  """Get the session list of chats in the workspace"""
-  chats: [String!]!
-
   """Get the context list of a session"""
   contexts(contextId: String, sessionId: String!): [CopilotContext!]!
+  docId: ID
   histories(docId: String, options: QueryChatHistoriesInput): [CopilotHistories!]!
 
   """Get the quota of the user in the workspace"""
   quota: CopilotQuota!
+
+  """Get the session list in the workspace"""
+  sessionIds(docId: String, options: QueryChatSessionsInput): [String!]!
   workspaceId: ID
 }
 
@@ -884,6 +882,10 @@ input QueryChatHistoriesInput {
   withPrompt: Boolean
 }
 
+input QueryChatSessionsInput {
+  action: Boolean
+}
+
 type QueryTooLongDataType {
   max: Int!
 }

packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts

@@ -301,6 +301,11 @@ declare global {
         docId: string,
         promptName?: string
       ) => Promise<string>;
+      getSessionIds: (
+        workspaceId: string,
+        docId?: string,
+        options?: { action?: boolean }
+      ) => Promise<string[] | undefined>;
       updateSession: (sessionId: string, promptName: string) => Promise<string>;
     }
 

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts

@@ -136,15 +136,23 @@ export class CopilotClient {
     }
   }
 
-  async getSessions(workspaceId: string) {
+  async getSessionIds(
+    workspaceId: string,
+    docId?: string,
+    options?: RequestOptions<
+      typeof getCopilotSessionsQuery
+    >['variables']['options']
+  ) {
     try {
       const res = await this.gql({
         query: getCopilotSessionsQuery,
         variables: {
           workspaceId,
+          docId,
+          options,
         },
       });
-      return res.currentUser?.copilot;
+      return res.currentUser?.copilot?.sessionIds;
     } catch (err) {
       throw resolveError(err);
     }

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts

@@ -1,5 +1,4 @@
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
-import type { ForkChatSessionInput } from '@affine/graphql';
 import { assertExists } from '@blocksuite/affine/global/utils';
 import { partition } from 'lodash-es';
 
@@ -48,36 +47,13 @@ export async function createChatSession({
     promptName,
   });
   // always update the prompt name
-  await updateChatSession({
+  await client.updateSession({
     sessionId,
-    client,
     promptName,
   });
   return sessionId;
 }
 
-export function updateChatSession({
-  client,
-  sessionId,
-  promptName,
-}: {
-  client: CopilotClient;
-  sessionId: string;
-  promptName: string;
-}) {
-  return client.updateSession({
-    sessionId,
-    promptName,
-  });
-}
-
-export function forkCopilotSession(
-  client: CopilotClient,
-  forkChatSessionInput: ForkChatSessionInput
-) {
-  return client.forkSession(forkChatSessionInput);
-}
-
 async function resizeImage(blob: Blob | File): Promise<Blob | null> {
   let src = '';
   try {
@@ -360,17 +336,3 @@ export function toImage({
     },
   };
 }
-
-export function cleanupSessions({
-  workspaceId,
-  docId,
-  sessionIds,
-  client,
-}: {
-  workspaceId: string;
-  docId: string;
-  sessionIds: string[];
-  client: CopilotClient;
-}) {
-  return client.cleanupSessions({ workspaceId, docId, sessionIds });
-}

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx

@@ -11,14 +11,7 @@ import { z } from 'zod';
 
 import type { CopilotClient } from './copilot-client';
 import type { PromptKey } from './prompt';
-import {
-  cleanupSessions,
-  createChatSession,
-  forkCopilotSession,
-  textToText,
-  toImage,
-  updateChatSession,
-} from './request';
+import { createChatSession, textToText, toImage } from './request';
 import { setupTracker } from './tracker';
 
 const filterStyleToPromptName = new Map(
@@ -424,9 +417,15 @@ Could you make a new website based on these notes and send back just the html fi
         promptName,
       });
     },
+    getSessionIds: async (
+      workspaceId: string,
+      docId?: string,
+      options?: { action?: boolean }
+    ) => {
+      return client.getSessionIds(workspaceId, docId, options);
+    },
     updateSession: async (sessionId: string, promptName: string) => {
-      return updateChatSession({
-        client,
+      return client.updateSession({
         sessionId,
         promptName,
       });
@@ -490,7 +489,7 @@ Could you make a new website based on these notes and send back just the html fi
       docId: string,
       sessionIds: string[]
     ) => {
-      await cleanupSessions({ workspaceId, docId, sessionIds, client });
+      await client.cleanupSessions({ workspaceId, docId, sessionIds });
     },
     ids: async (
       workspaceId: string,
@@ -533,7 +532,7 @@ Could you make a new website based on these notes and send back just the html fi
   AIProvider.provide('onboarding', toggleGeneralAIOnboarding);
 
   AIProvider.provide('forkChat', options => {
-    return forkCopilotSession(client, options);
+    return client.forkSession(options);
   });
 
   const disposeRequestLoginHandler = AIProvider.slots.requestLogin.on(() => {

packages/frontend/graphql/src/graphql/copilot-sessions-get.gql

@@ -1,8 +1,11 @@
-query getCopilotSessions($workspaceId: String!) {
+query getCopilotSessions(
+  $workspaceId: String!
+  $docId: String
+  $options: QueryChatSessionsInput
+) {
   currentUser {
     copilot(workspaceId: $workspaceId) {
-      actions
-      chats
+      sessionIds(docId: $docId, options: $options)
     }
   }
 }

packages/frontend/graphql/src/graphql/index.ts

@@ -421,11 +421,10 @@ export const getCopilotSessionsQuery = {
   definitionName: 'currentUser',
   containsFile: false,
   query: `
-query getCopilotSessions($workspaceId: String!) {
+query getCopilotSessions($workspaceId: String!, $docId: String, $options: QueryChatSessionsInput) {
   currentUser {
     copilot(workspaceId: $workspaceId) {
-      actions
-      chats
+      sessionIds(docId: $docId, options: $options)
     }
   }
 }`,

packages/frontend/graphql/src/schema.ts

@@ -76,15 +76,14 @@ export enum ContextFileStatus {
 
 export interface Copilot {
   __typename?: 'Copilot';
-  /** Get the session list of actions in the workspace */
-  actions: Array<Scalars['String']['output']>;
-  /** Get the session list of chats in the workspace */
-  chats: Array<Scalars['String']['output']>;
   /** Get the context list of a session */
   contexts: Array<CopilotContext>;
+  docId: Maybe<Scalars['ID']['output']>;
   histories: Array<CopilotHistories>;
   /** Get the quota of the user in the workspace */
   quota: CopilotQuota;
+  /** Get the session list in the workspace */
+  sessionIds: Array<Scalars['String']['output']>;
   workspaceId: Maybe<Scalars['ID']['output']>;
 }
 
@@ -98,6 +97,11 @@ export interface CopilotHistoriesArgs {
   options?: InputMaybe<QueryChatHistoriesInput>;
 }
 
+export interface CopilotSessionIdsArgs {
+  docId?: InputMaybe<Scalars['String']['input']>;
+  options?: InputMaybe<QueryChatSessionsInput>;
+}
+
 export interface CopilotContext {
   __typename?: 'CopilotContext';
   /** list files in context */
@@ -1301,6 +1305,10 @@ export interface QueryChatHistoriesInput {
   withPrompt?: InputMaybe<Scalars['Boolean']['input']>;
 }
 
+export interface QueryChatSessionsInput {
+  action?: InputMaybe<Scalars['Boolean']['input']>;
+}
+
 export interface QueryTooLongDataType {
   __typename?: 'QueryTooLongDataType';
   max: Scalars['Int']['output'];
@@ -2196,17 +2204,15 @@ export type UpdateCopilotSessionMutation = {
 
 export type GetCopilotSessionsQueryVariables = Exact<{
   workspaceId: Scalars['String']['input'];
+  docId?: InputMaybe<Scalars['String']['input']>;
+  options?: InputMaybe<QueryChatSessionsInput>;
 }>;
 
 export type GetCopilotSessionsQuery = {
   __typename?: 'Query';
   currentUser: {
     __typename?: 'UserType';
-    copilot: {
-      __typename?: 'Copilot';
-      actions: Array<string>;
-      chats: Array<string>;
-    };
+    copilot: { __typename?: 'Copilot'; sessionIds: Array<string> };
   } | null;
 };