From e4b816f1535368b9e52b259fe38c02b269860712 Mon Sep 17 00:00:00 2001
From: darkskygit <darksky.graphite@toeverything.info>
Date: Mon, 15 Jul 2024 04:16:21 +0000
Subject: [PATCH] feat: add ping for event source (#7493)

---
 .../server/src/plugins/copilot/controller.ts  | 33 ++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts
index ec46455d1..6198ba849 100644
--- a/packages/backend/server/src/plugins/copilot/controller.ts
+++ b/packages/backend/server/src/plugins/copilot/controller.ts
@@ -14,12 +14,16 @@ import {
   concatMap,
   connect,
   EMPTY,
+  finalize,
   from,
+  interval,
   map,
   merge,
   mergeMap,
   Observable,
+  Subject,
   switchMap,
+  takeUntil,
   toArray,
 } from 'rxjs';
 
@@ -41,7 +45,7 @@ import { CopilotCapability, CopilotTextProvider } from './types';
 import { CopilotWorkflowService, GraphExecutorState } from './workflow';
 
 export interface ChatEvent {
-  type: 'event' | 'attachment' | 'message' | 'error';
+  type: 'event' | 'attachment' | 'message' | 'error' | 'ping';
   id?: string;
   data: string | object;
 }
@@ -51,6 +55,8 @@ type CheckResult = {
   hasAttachment?: boolean;
 };
 
+const PING_INTERVAL = 5000;
+
 @Controller('/api/copilot')
 export class CopilotController {
   private readonly logger = new Logger(CopilotController.name);
@@ -159,6 +165,19 @@ export class CopilotController {
     return num;
   }
 
+  private mergePingStream(
+    messageId: string,
+    source$: Observable<ChatEvent>
+  ): Observable<ChatEvent> {
+    const subject$ = new Subject();
+    const ping$ = interval(PING_INTERVAL).pipe(
+      map(() => ({ type: 'ping' as const, id: messageId, data: '' })),
+      takeUntil(subject$)
+    );
+
+    return merge(source$.pipe(finalize(() => subject$.next(null))), ping$);
+  }
+
   @Get('/chat/:sessionId')
   async chat(
     @CurrentUser() user: CurrentUser,
@@ -216,7 +235,7 @@ export class CopilotController {
 
       const session = await this.appendSessionMessage(sessionId, messageId);
 
-      return from(
+      const source$ = from(
         provider.generateTextStream(session.finish(params), session.model, {
           ...session.config.promptConfig,
           signal: this.getSignal(req),
@@ -246,6 +265,8 @@ export class CopilotController {
         ),
         catchError(mapSseError)
       );
+
+      return this.mergePingStream(messageId, source$);
     } catch (err) {
       return mapSseError(err);
     }
@@ -270,7 +291,7 @@ export class CopilotController {
         });
       }
 
-      return from(
+      const source$ = from(
         this.workflow.runGraph(params, session.model, {
           ...session.config.promptConfig,
           signal: this.getSignal(req),
@@ -316,6 +337,8 @@ export class CopilotController {
         ),
         catchError(mapSseError)
       );
+
+      return this.mergePingStream(messageId, source$);
     } catch (err) {
       return mapSseError(err);
     }
@@ -353,7 +376,7 @@ export class CopilotController {
         sessionId
       );
 
-      return from(
+      const source$ = from(
         provider.generateImagesStream(session.finish(params), session.model, {
           seed: this.parseNumber(params.seed),
           signal: this.getSignal(req),
@@ -389,6 +412,8 @@ export class CopilotController {
         ),
         catchError(mapSseError)
       );
+
+      return this.mergePingStream(messageId, source$);
     } catch (err) {
       return mapSseError(err);
     }