From eb49ffaedbf36177b1a3f3c99b7f7c06e8232330 Mon Sep 17 00:00:00 2001
From: akumatus <akumatus@gmail.com>
Date: Wed, 28 May 2025 07:34:23 +0000
Subject: [PATCH] feat(core): support fork session without latestMessageId
 (#12587)

Close [AI-86](https://linear.app/affine-design/issue/AI-86)

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->

## Summary by CodeRabbit

- **New Features**
  - Improved chat session forking to allow creating a fork without specifying the latest message, enabling more flexible session management.

- **Bug Fixes**
  - Forking a chat session with an invalid latest message ID now correctly returns an error.

- **Tests**
  - Added and updated test cases to cover session forking with missing or invalid latest message IDs, ensuring robust behavior in these scenarios.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 .../__snapshots__/copilot.spec.ts.md          |  32 +++++++++++++++++
 .../__snapshots__/copilot.spec.ts.snap        | Bin 1429 -> 1440 bytes
 .../server/src/__tests__/copilot.e2e.ts       |  23 +++++++++++-
 .../server/src/__tests__/copilot.spec.ts      |  33 ++++++++++++++++--
 .../server/src/__tests__/utils/copilot.ts     |   2 +-
 .../server/src/plugins/copilot/resolver.ts    |   3 +-
 .../server/src/plugins/copilot/session.ts     |  21 ++++++-----
 .../server/src/plugins/copilot/types.ts       |   2 +-
 packages/backend/server/src/schema.gql        |   2 +-
 packages/common/graphql/src/schema.ts         |   2 +-
 .../core/src/blocksuite/ai/actions/types.ts   |   2 +-
 11 files changed, 103 insertions(+), 19 deletions(-)

diff --git a/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md b/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md
index ee137de87d..39a6e2be9a 100644
--- a/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md
+++ b/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md
@@ -116,6 +116,38 @@ Generated by [AVA](https://avajs.dev).
       },
     ]
 
+> should generate the final message
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: 'hello',
+        params: {},
+        role: 'user',
+      },
+      {
+        content: 'world',
+        params: {},
+        role: 'assistant',
+      },
+      {
+        content: 'aaa',
+        params: {},
+        role: 'user',
+      },
+      {
+        content: 'bbb',
+        params: {},
+        role: 'assistant',
+      },
+    ]
+
 ## should revert message correctly
 
 > should have three messages before revert
diff --git a/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.snap b/packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.snap
index e2b001246740af92fd44403957aabd24f88585cf..943a031e2df0991da530985188df524f3006b04d 100644
GIT binary patch
literal 1440
zcmV;R1z-9>RzV<BZeeh9Xm4~Nb~6eC07}aw^(1lK6w^XE-hzM?Oa6{561^iaYnsV0
zH8E^MdmoDk00000000B+nqO!fMI6V!zuCP@nrrjNivK8ft%^Zn(&mr0C8!leYn37u
zqy-CecXNN5z1w4V@=rlfQ=uyKK@o+%G@yd*gMy$4;zKPIK?TM7R20RE5B9}A6%=px
zcK0rOF_XJAX)BwD%O*QJznz)S_dCCto8!5(X=Nsyxp^Ur(iAgyIt6Re%o=H7@U$rm
zX&EJ6=6PXc3S1hFaGau5E;?miaSE1nUYgHlx#aulq=pW-Kh#OuRsi<`co4uwO7>A=
zmu>S|qTs5hrw7ypVxSb!p4x*cqUYS-<kq3T3*ZERvj8px^(6PE4~dK<7GC!D_Ie@v
z2)L1eaf%ex7Z%+B(m*}fNbGv|D$z$B8LKQsSq6Wn3c@rkW7@LKEU264SptHa$qKi5
z$%#~<+m<Q9BF?NMMahfB6tNOstfgwxCmmtCO=NQfyc%tu*3>+T);4X~*)X2hDq-@p
zFlTP_2Rkpy!sb#KvLK9HvCK`QBpiq5MLfFK?&yw1M>ku{<%BKDG8kQBx+n_<FI%!8
zY@^onzzA+{HiDdGPyB}s%=w_O6Ahr~qX2dSco@K56)c*E$bk@+BR(uCqEFpR=2iJS
zfIk7;M8MW)Z#8kF?;~J80WT2nmV%v5(mx^Ka{_)L;8L`Ba^<KQtZmd(u>#b(i5_h&
zijUxOV$@qAT8pucJ5JG&yevJ|UuM7&wlwRV3-2bDyxXxBE*Qfd!~MU7TOVkEJ2?*E
zAplPRcm_baN@hs~dBWp)CWQR`I^<`4$WufKLZF9H_fq%Z%n(qlPQ-5#a8ix_w*>s)
zb9`O2uba^yVZaUs9Av;$w0N?q_%Q~|GvG7>&P9uNHWmMe0lEg<r2!9Uk>Xt|7nac0
zLRhc27Nsp=eb8ErMqw=@kj6Rp!kk&HrfBJO+KX_${f1OiJzWR7^wo2S$lY~I-c7~s
z61z+6F0s2@i|*ntRyX@4Z80V6`tmko*|x|?bGE)veZ@B<ZTo5((sy-3`WvAk_16t4
zNZs8<|1qT@X<ra<nt=1lEMMzjzy<~kMx7-(#-u&XfCCJ8Ug=p{;8~}lMVoWemJR5S
z_SX%n{TJG{JY_2D8E(Nyl$DKKrk5IR!4!Hm;D+ezv69`TeqVA0KIID279x<HW28mS
zvW2h4VDbD>gWwX|0o)JZF#!7j98`%-_Y_nN#rAexY^UmCyWoqhlk}Sj*i69V)fI1A
zt+2xcytY_&(5$ep3HY9Xi_xOZ3cG~??y`87LSgO5CLPOW(S1~R>3vkO{BDW#fpBsu
zSyRFY-#VOTcuopCLi^BaY!!W6-zxeEz;6J03Fuejy@!CO2(Z<7zd^t|1bn8zyB%hK
zUcIgMwP&`ISTNaMQP$#ACQfCpW>1b&nK+e+Q<*rGiBp-Ymf0(x%J>^OcQl_muqVq#
zvE)7=W<+LETJ|ERk0u)!Wdi_)RczKU91jn*<o!J9u;MUEt>3$ejzbB6oUPQlTd$$c
z&H(r(8eq{*?k4)RA%Hd^X1&*3?`}y~Q>aI1Vf*w&&fctdUnZ<#OF_ztvvX)j%<xjh
z6oVP7G$e&1@9}zf@K7c1KSC5dbF1X}Z3Bkn`C;#K#QPleKF0<Of221>_vvcG<w#rb
zQs5#_UN#qbvF;-8g)YLr1n`yeAQ8I=;8G-_wjg55U%qyftgJA1hbri*cWMJ^It7g%
zAz%jq(*zt{`FSSA7<njt*qT|@XP_SmIHz{c8nAygMA0SBJ07u{s`>U72HeLM_IgrD
utXs!u6O-<L$9(ec4n3xv6}GoOT(PZEMb=I>)mmzBH_E>!i_)~}C;$MnS+vLi

literal 1429
zcmV;G1#0?1RzV<BZeeh9Xm4~Nb~6eC03`-mp+>4JrX$|XA}#SbE2<S_fki6B9u<0!
zzK^Af+aHSv00000000B+nqO!fMI6V!zuCRZHP_~k75`D}S`~xDq|F~~OHeC{)+$9R
zNDCI`?&kh9ySK;g<e!3|ZH21P2SpV6(t--M4+?@Jh!3?;1QitPQ&AKvKG+xgR8YL#
z+ughD#Y`^grByBun@x6hemgUt?{|JXv&Zup)6P!1hfWAtlBSrw(=FOlX3oe6gJ(=(
zNZYV@nHPkSEplnN!gWh_x#X64#Vy*>ed$Cl$0grOr!;iI`=JzRn*rPl;6VWEDcMhn
z9gf52h{Chp-d<2Qh=I~XyJ|P4iJtR*lUIlS4uF#Y&H=a>)|1?wIUur<SajOg*XM`q
zC*Vc`CMZ@^pI>wxNCWj^J+bS(vqV32W$m&QWf}gRE(+7MjTzf9bD$ogX9)-|CM(?G
zmK&==cWhHcMcg@83d@hh6tNP1ti@{6r(EH9O=O1%cqQIEt*LnwuWiP5a#1|LRl?*M
zVb0zb40b`3g~O#VWKkITQkk2EC0v&mL}zsWyQ4cEAKhFjpBIiO%W!m!nUX9Tyll&&
zaEw~hLnFAo*$DEsGx;wza4-OcoooO_9|N!*z{3D`t6<T5MD|6n90_1a6Mf>HGQY~-
z0Q>>qCIU9cd#i~XeGdV933!2kHx=xpNdK6C&j|RLfXngX$(5sKuw|pBiWQ>PL-c5C
zQ33>)5~JP{(OQf(+;vN?<Ynozevko2*y60G=3h-LdbMjcT(ARo2k!qZ-1=Yx+{p<5
z4*_@rz%u|cRWdCV<O!eW*$DFY>X4rcAWst|2%#Rvyi?tWGfP0JIuXA?z$rEQ-w^P9
z!0~nQzHUZ;lmXiqu%7|b@#4v*;>Q?pf&ph3a6Vo<)l~d%2Iv}amj*nb#fo>YTv$R|
z4Pm|3T9jo8>;2YZGzx18fi%v!ALeYknxbVg89%~><u|07>ghVrt+&r1VprEJdNti%
zcL~zTjX_CAObe%;fM#vS5m{-@)l<or14Gibw$qTltsBx`hzx0<Zb;$e>ZRq!l!l~z
zPQV!gE-2?=t&0Kc7%&ufmiQQx_A~?bG2nTnXKA5losJi6o;_MNpx>82v#Z{aXv@-+
zsjO$D1tU>bHhP5~YODoQ=+l53;<Lv}-v0&y$rbsuCrn3(P<F175qaAYff~c~_QM9j
zCAI>%AHZV(_5j$g5}WQTs1%9qt-9Dw*Tr@*5L=4$n+e!Rz~j}NG^1A7VFF%VC_89Y
z*jEI6N5G|c(Po9+!T>LE-Jwv}a%7X~$Y#Nn4lg-Asz{JqB7HELTr7K97|~Nl(+tl`
z;lyYkZpX{qNA;Jv9|8OdppSq7HQu`jc!~f=jrZ#WyiLHT3cQ!Y>`(1m3|oC>Q^dl_
z_NubhNo6{z%(d*vom8fi%5+khPAbz$Wv*FfuYM{MEacqLyf5fXm5q|+-5_Q~c1qgL
z0^bKrH89Es0gR~FtYJ72Z95l&Jn67vySUcxJw(T$1VGMJYTd2ZP-kZWd>s$4U?=ww
z{n8LXn-p{YW3Kg<q^l{^BQ(E#`U1P-YrThws@PPNvf^$Z9u~9Qs+eLZYg@xoxbhyq
zcZUvC3c)Qz;XSuXUf42dNM0E6kE8x^%s-9~8o@|!h_BPthRc<X;8y4&PhK$>d9m&y
z??x`dz5wv0a;u+R0&qDNQJWXB<?WF!%gzaNXQYDedZ#v!rc>DXQ3AFRFhju6m49)d
z7$YC04_Y&;`V8~~0q523Sp)XZhA6t^1=lBbLp9&t#DM$Q{9aEgiS_CjYhu#-?|4Ao
j-I2$XbHef0hbxY4Rb=hkgIY@s?nU_rGR#`cj3fX6e+#oo

diff --git a/packages/backend/server/src/__tests__/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot.e2e.ts
index 3ee870cb78..3e5537a0dd 100644
--- a/packages/backend/server/src/__tests__/copilot.e2e.ts
+++ b/packages/backend/server/src/__tests__/copilot.e2e.ts
@@ -281,7 +281,7 @@ test('should fork session correctly', async t => {
   const assertForkSession = async (
     workspaceId: string,
     sessionId: string,
-    lastMessageId: string,
+    lastMessageId: string | undefined,
     error: string,
     asserter = async (x: any) => {
       const forkedSessionId = await x;
@@ -330,6 +330,27 @@ test('should fork session correctly', async t => {
     );
   }
 
+  // should be able to fork session without latestMessageId (copy all messages)
+  {
+    forkedSessionId = await assertForkSession(
+      id,
+      sessionId,
+      undefined,
+      'should be able to fork session without latestMessageId'
+    );
+  }
+
+  // should not be able to fork session with wrong latestMessageId
+  {
+    await assertForkSession(id, sessionId, 'wrong-message-id', '', async x => {
+      await t.throwsAsync(
+        x,
+        { instanceOf: Error },
+        'should not able to fork session with wrong latestMessageId'
+      );
+    });
+  }
+
   {
     const u2 = await app.signupV1('u2@affine.pro');
     await assertForkSession(id, sessionId, randomUUID(), '', async x => {
diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts
index 233bd1584e..547878928a 100644
--- a/packages/backend/server/src/__tests__/copilot.spec.ts
+++ b/packages/backend/server/src/__tests__/copilot.spec.ts
@@ -410,6 +410,27 @@ test('should be able to fork chat session', async t => {
     'should fork new session with same params'
   );
 
+  // fork session without latestMessageId
+  const forkedSessionId3 = await session.fork({
+    userId,
+    sessionId,
+    ...commonParams,
+  });
+
+  // fork session with wrong latestMessageId
+  await t.throwsAsync(
+    session.fork({
+      userId,
+      sessionId,
+      latestMessageId: 'wrong-message-id',
+      ...commonParams,
+    }),
+    {
+      instanceOf: Error,
+    },
+    'should not able to fork new session with wrong latestMessageId'
+  );
+
   const cleanObject = (obj: any[]) =>
     JSON.parse(
       JSON.stringify(obj, (k, v) =>
@@ -436,11 +457,17 @@ test('should be able to fork chat session', async t => {
     t.snapshot(cleanObject(finalMessages), 'should generate the final message');
   }
 
+  // check third times forked session
+  {
+    const s3 = (await session.get(forkedSessionId3))!;
+    const finalMessages = s3.finish(params);
+    t.snapshot(cleanObject(finalMessages), 'should generate the final message');
+  }
+
   // check original session messages
   {
-    const s3 = (await session.get(sessionId))!;
-
-    const finalMessages = s3.finish(params);
+    const s4 = (await session.get(sessionId))!;
+    const finalMessages = s4.finish(params);
     t.snapshot(cleanObject(finalMessages), 'should generate the final message');
   }
 
diff --git a/packages/backend/server/src/__tests__/utils/copilot.ts b/packages/backend/server/src/__tests__/utils/copilot.ts
index 59f1af7b69..8ae7bb6a67 100644
--- a/packages/backend/server/src/__tests__/utils/copilot.ts
+++ b/packages/backend/server/src/__tests__/utils/copilot.ts
@@ -57,7 +57,7 @@ export async function forkCopilotSession(
   workspaceId: string,
   docId: string,
   sessionId: string,
-  latestMessageId: string
+  latestMessageId?: string
 ): Promise<string> {
   const res = await app.gql(
     `
diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts
index 53214f6559..b70634f8c4 100644
--- a/packages/backend/server/src/plugins/copilot/resolver.ts
+++ b/packages/backend/server/src/plugins/copilot/resolver.ts
@@ -91,8 +91,9 @@ class ForkChatSessionInput {
   @Field(() => String, {
     description:
       'Identify a message in the array and keep it with all previous messages into a forked session.',
+    nullable: true,
   })
-  latestMessageId!: string;
+  latestMessageId?: string;
 }
 
 @InputType()
diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts
index e160098b27..c2dadeb31b 100644
--- a/packages/backend/server/src/plugins/copilot/session.ts
+++ b/packages/backend/server/src/plugins/copilot/session.ts
@@ -673,16 +673,19 @@ export class ChatSessionService {
     if (!state) {
       throw new CopilotSessionNotFound();
     }
-    const lastMessageIdx = state.messages.findLastIndex(
-      ({ id, role }) =>
-        role === AiPromptRole.assistant && id === options.latestMessageId
-    );
-    if (lastMessageIdx < 0) {
-      throw new CopilotMessageNotFound({ messageId: options.latestMessageId });
+    let messages = state.messages.map(m => ({ ...m, id: undefined }));
+    if (options.latestMessageId) {
+      const lastMessageIdx = state.messages.findLastIndex(
+        ({ id, role }) =>
+          role === AiPromptRole.assistant && id === options.latestMessageId
+      );
+      if (lastMessageIdx < 0) {
+        throw new CopilotMessageNotFound({
+          messageId: options.latestMessageId,
+        });
+      }
+      messages = messages.slice(0, lastMessageIdx + 1);
     }
-    const messages = state.messages
-      .slice(0, lastMessageIdx + 1)
-      .map(m => ({ ...m, id: undefined }));
 
     const forkedState = {
       ...state,
diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts
index ae403bf43d..9aefd70f55 100644
--- a/packages/backend/server/src/plugins/copilot/types.ts
+++ b/packages/backend/server/src/plugins/copilot/types.ts
@@ -119,7 +119,7 @@ export interface ChatSessionPromptUpdateOptions
 export interface ChatSessionForkOptions
   extends Omit<ChatSessionOptions, 'promptName'> {
   sessionId: string;
-  latestMessageId: string;
+  latestMessageId?: string;
 }
 
 export interface ChatSessionState
diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql
index ff30e3190e..f9b8da63a4 100644
--- a/packages/backend/server/src/schema.gql
+++ b/packages/backend/server/src/schema.gql
@@ -703,7 +703,7 @@ input ForkChatSessionInput {
   """
   Identify a message in the array and keep it with all previous messages into a forked session.
   """
-  latestMessageId: String!
+  latestMessageId: String
   sessionId: String!
   workspaceId: String!
 }
diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts
index 1725f57530..ccebb1b72d 100644
--- a/packages/common/graphql/src/schema.ts
+++ b/packages/common/graphql/src/schema.ts
@@ -879,7 +879,7 @@ export enum FeatureType {
 export interface ForkChatSessionInput {
   docId: Scalars['String']['input'];
   /** Identify a message in the array and keep it with all previous messages into a forked session. */
-  latestMessageId: Scalars['String']['input'];
+  latestMessageId?: InputMaybe<Scalars['String']['input']>;
   sessionId: Scalars['String']['input'];
   workspaceId: Scalars['String']['input'];
 }
diff --git a/packages/frontend/core/src/blocksuite/ai/actions/types.ts b/packages/frontend/core/src/blocksuite/ai/actions/types.ts
index 40e701c427..f690489d9d 100644
--- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts
+++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts
@@ -93,7 +93,7 @@ declare global {
       docId: string;
       workspaceId: string;
       sessionId: string;
-      latestMessageId: string;
+      latestMessageId?: string;
     }
 
     interface AIImageActionOptions extends AITextActionOptions {