feat: allow retry with new message (#10307)

fix AF-1630

packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md

@@ -4,6 +4,36 @@ The actual snapshot is saved in `copilot.e2e.ts.snap`.
 
 Generated by [AVA](https://avajs.dev).
 
+## should be able to retry with api
+
+> should be able to list history after retry
+
+    [
+      {
+        messages: [
+          {
+            content: 'generate text to text',
+            role: 'assistant',
+          },
+        ],
+        tokens: 0,
+      },
+    ]
+
+> should be able to list history after retry
+
+    [
+      {
+        messages: [
+          {
+            content: 'generate text to text',
+            role: 'assistant',
+          },
+        ],
+        tokens: 0,
+      },
+    ]
+
 ## should be able to manage context
 
 > should list context files
@@ -13,14 +43,3 @@ Generated by [AVA](https://avajs.dev).
         id: 'docId1',
       },
     ]
-
-> should list context docs
-
-    [
-      {
-        blobId: 'fileId1',
-        chunkSize: 3,
-        name: 'sample.pdf',
-        status: 'finished',
-      },
-    ]

packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md

@@ -0,0 +1,151 @@
+# Snapshot report for `src/__tests__/copilot.spec.ts`
+
+The actual snapshot is saved in `copilot.spec.ts.snap`.
+
+Generated by [AVA](https://avajs.dev).
+
+## should revert message correctly
+
+> should have three messages before revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+      {
+        content: '3',
+        role: 'user',
+      },
+      {
+        content: '4',
+        role: 'assistant',
+      },
+    ]
+
+> should remove assistant message after revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+      {
+        content: '3',
+        role: 'user',
+      },
+    ]
+
+> should remove assistant message after revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+    ]
+
+> should have three messages before revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+      {
+        content: '3',
+        role: 'user',
+      },
+      {
+        content: '4',
+        role: 'assistant',
+      },
+    ]
+
+> should remove assistant message after revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+      {
+        content: '3',
+        role: 'user',
+      },
+    ]
+
+> should remove assistant message after revert
+
+    [
+      {
+        content: 'hello world',
+        params: {
+          word: 'world',
+        },
+        role: 'system',
+      },
+      {
+        content: '1',
+        role: 'user',
+      },
+      {
+        content: '2',
+        role: 'assistant',
+      },
+    ]

packages/backend/server/src/__tests__/copilot.e2e.ts

@@ -498,6 +498,15 @@ test('should be able to retry with api', async t => {
     );
   }
 
+  const cleanObject = (obj: any[]) =>
+    JSON.parse(
+      JSON.stringify(obj, (k, v) =>
+        ['id', 'sessionId', 'createdAt'].includes(k) || v === null
+          ? undefined
+          : v
+      )
+    );
+
   // retry chat
   {
     const { id } = await createWorkspace(app);
@@ -514,10 +523,32 @@ test('should be able to retry with api', async t => {
 
     // should only have 1 message
     const histories = await getHistories(app, { workspaceId: id });
-    t.deepEqual(
-      histories.map(h => h.messages.map(m => m.content)),
-      [['generate text to text']],
-      'should be able to list history'
+    t.snapshot(
+      cleanObject(histories),
+      'should be able to list history after retry'
+    );
+  }
+
+  // retry chat with new message id
+  {
+    const { id } = await createWorkspace(app);
+    const sessionId = await createCopilotSession(
+      app,
+      id,
+      randomUUID(),
+      promptName
+    );
+    const messageId = await createCopilotMessage(app, sessionId);
+    await chatWithText(app, sessionId, messageId);
+    // retry with new message id
+    const newMessageId = await createCopilotMessage(app, sessionId);
+    await chatWithText(app, sessionId, newMessageId, '', true);
+
+    // should only have 1 message
+    const histories = await getHistories(app, { workspaceId: id });
+    t.snapshot(
+      cleanObject(histories),
+      'should be able to list history after retry'
     );
   }
 

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -602,36 +602,81 @@ test('should revert message correctly', async t => {
 
     const message = (await session.createMessage({
       sessionId,
-      content: 'hello',
+      content: '1',
     }))!;
 
     await s.pushByMessageId(message);
     await s.save();
   }
 
+  const cleanObject = (obj: any[]) =>
+    JSON.parse(
+      JSON.stringify(obj, (k, v) =>
+        ['id', 'createdAt'].includes(k) || v === null ? undefined : v
+      )
+    );
+
   // check ChatSession behavior
   {
     const s = (await session.get(sessionId))!;
-    s.push({ role: 'assistant', content: 'hi', createdAt: new Date() });
+    s.push({ role: 'assistant', content: '2', createdAt: new Date() });
+    s.push({ role: 'user', content: '3', createdAt: new Date() });
+    s.push({ role: 'assistant', content: '4', createdAt: new Date() });
     await s.save();
     const beforeRevert = s.finish({ word: 'world' });
-    t.is(beforeRevert.length, 3, 'should have three messages before revert');
+    t.snapshot(
+      cleanObject(beforeRevert),
+      'should have three messages before revert'
+    );
 
-    s.revertLatestMessage();
-    const afterRevert = s.finish({ word: 'world' });
-    t.is(afterRevert.length, 2, 'should remove assistant message after revert');
+    {
+      s.revertLatestMessage(false);
+      const afterRevert = s.finish({ word: 'world' });
+      t.snapshot(
+        cleanObject(afterRevert),
+        'should remove assistant message after revert'
+      );
+    }
+
+    {
+      s.revertLatestMessage(true);
+      const afterRevert = s.finish({ word: 'world' });
+      t.snapshot(
+        cleanObject(afterRevert),
+        'should remove assistant message after revert'
+      );
+    }
   }
 
   // check database behavior
   {
     let s = (await session.get(sessionId))!;
-    const beforeRevert = s.finish({ word: 'world' });
-    t.is(beforeRevert.length, 3, 'should have three messages before revert');
 
-    await session.revertLatestMessage(sessionId);
-    s = (await session.get(sessionId))!;
-    const afterRevert = s.finish({ word: 'world' });
-    t.is(afterRevert.length, 2, 'should remove assistant message after revert');
+    const beforeRevert = s.finish({ word: 'world' });
+    t.snapshot(
+      cleanObject(beforeRevert),
+      'should have three messages before revert'
+    );
+
+    {
+      await session.revertLatestMessage(sessionId, false);
+      s = (await session.get(sessionId))!;
+      const afterRevert = s.finish({ word: 'world' });
+      t.snapshot(
+        cleanObject(afterRevert),
+        'should remove assistant message after revert'
+      );
+    }
+
+    {
+      await session.revertLatestMessage(sessionId, true);
+      s = (await session.get(sessionId))!;
+      const afterRevert = s.finish({ word: 'world' });
+      t.snapshot(
+        cleanObject(afterRevert),
+        'should remove assistant message after revert'
+      );
+    }
   }
 });
 

packages/backend/server/src/__tests__/utils/copilot.ts

@@ -444,9 +444,12 @@ export async function chatWithText(
   app: TestingApp,
   sessionId: string,
   messageId?: string,
-  prefix = ''
+  prefix = '',
+  retry?: boolean
 ): Promise<string> {
-  const query = messageId ? `?messageId=${messageId}` : '';
+  const query = messageId
+    ? `?messageId=${messageId}` + (retry ? '&retry=true' : '')
+    : '';
   const res = await app
     .GET(`/api/copilot/chat/${sessionId}${prefix}${query}`)
     .expect(200);