feat: add retry support for copilot (#6947)

packages/backend/server/tests/copilot.e2e.ts

@@ -228,6 +228,61 @@ test('should be able to chat with api', async t => {
   Sinon.restore();
 });
 
+test('should be able to retry with api', async t => {
+  const { app, storage } = t.context;
+
+  Sinon.stub(storage, 'handleRemoteLink').resolvesArg(2);
+
+  // normal chat
+  {
+    const { id } = await createWorkspace(app, token);
+    const sessionId = await createCopilotSession(
+      app,
+      token,
+      id,
+      randomUUID(),
+      promptName
+    );
+    const messageId = await createCopilotMessage(app, token, sessionId);
+    // chat 2 times
+    await chatWithText(app, token, sessionId, messageId);
+    await chatWithText(app, token, sessionId, messageId);
+
+    const histories = await getHistories(app, token, { workspaceId: id });
+    t.deepEqual(
+      histories.map(h => h.messages.map(m => m.content)),
+      [['generate text to text', 'generate text to text']],
+      'should be able to list history'
+    );
+  }
+
+  // retry chat
+  {
+    const { id } = await createWorkspace(app, token);
+    const sessionId = await createCopilotSession(
+      app,
+      token,
+      id,
+      randomUUID(),
+      promptName
+    );
+    const messageId = await createCopilotMessage(app, token, sessionId);
+    await chatWithText(app, token, sessionId, messageId);
+    // retry without message id
+    await chatWithText(app, token, sessionId);
+
+    // should only have 1 message
+    const histories = await getHistories(app, token, { workspaceId: id });
+    t.deepEqual(
+      histories.map(h => h.messages.map(m => m.content)),
+      [['generate text to text']],
+      'should be able to list history'
+    );
+  }
+
+  Sinon.restore();
+});
+
 test('should reject message from different session', async t => {
   const { app } = t.context;
 

packages/backend/server/tests/copilot.spec.ts

@@ -362,6 +362,59 @@ test('should save message correctly', async t => {
   t.is(s.stashMessages.length, 0, 'should empty stash messages after save');
 });
 
+test('should revert message correctly', async t => {
+  const { prompt, session } = t.context;
+
+  // init session
+  let sessionId: string;
+  {
+    await prompt.set('prompt', 'model', [
+      { role: 'system', content: 'hello {{word}}' },
+    ]);
+
+    sessionId = await session.create({
+      docId: 'test',
+      workspaceId: 'test',
+      userId,
+      promptName: 'prompt',
+    });
+    const s = (await session.get(sessionId))!;
+
+    const message = (await session.createMessage({
+      sessionId,
+      content: 'hello',
+    }))!;
+
+    await s.pushByMessageId(message);
+    await s.save();
+  }
+
+  // check ChatSession behavior
+  {
+    const s = (await session.get(sessionId))!;
+    s.push({ role: 'assistant', content: 'hi', createdAt: new Date() });
+    await s.save();
+    const beforeRevert = s.finish({ word: 'world' });
+    t.is(beforeRevert.length, 3, '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');
+  }
+
+  // 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');
+  }
+});
+
 // ==================== provider ====================
 
 test('should be able to get provider', async t => {

packages/backend/server/tests/utils/copilot.ts

@@ -196,11 +196,12 @@ export async function chatWithText(
   app: INestApplication,
   userToken: string,
   sessionId: string,
-  messageId: string,
+  messageId?: string,
   prefix = ''
 ): Promise<string> {
+  const query = messageId ? `?messageId=${messageId}` : '';
   const res = await request(app.getHttpServer())
-    .get(`/api/copilot/chat/${sessionId}${prefix}?messageId=${messageId}`)
+    .get(`/api/copilot/chat/${sessionId}${prefix}${query}`)
     .auth(userToken, { type: 'bearer' })
     .expect(200);
 
@@ -211,7 +212,7 @@ export async function chatWithTextStream(
   app: INestApplication,
   userToken: string,
   sessionId: string,
-  messageId: string
+  messageId?: string
 ) {
   return chatWithText(app, userToken, sessionId, messageId, '/stream');
 }
@@ -220,7 +221,7 @@ export async function chatWithImages(
   app: INestApplication,
   userToken: string,
   sessionId: string,
-  messageId: string
+  messageId?: string
 ) {
   return chatWithText(app, userToken, sessionId, messageId, '/images');
 }