feat: handle stream error (#6653)

toeverything · Apr 23, 2024 · 64ad83f · 64ad83f
1 parent 01a0f60
commit 64ad83f
Showing 1 changed file with 45 additions and 3 deletions.
diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts
@@ -2,6 +2,7 @@ import {
   BadRequestException,
   Controller,
   Get,
+  HttpException,
   InternalServerErrorException,
   Logger,
   NotFoundException,
@@ -13,6 +14,7 @@ import {
 } from '@nestjs/common';
 import type { Request, Response } from 'express';
 import {
+  catchError,
   concatMap,
   connect,
   EMPTY,
@@ -21,6 +23,7 @@ import {
   merge,
   mergeMap,
   Observable,
+  of,
   switchMap,
   toArray,
 } from 'rxjs';
@@ -34,7 +37,7 @@ import { CopilotStorage } from './storage';
 import { CopilotCapability } from './types';
 
 export interface ChatEvent {
-  type: 'attachment' | 'message';
+  type: 'attachment' | 'message' | 'error';
   id?: string;
   data: string;
 }
@@ -83,6 +86,19 @@ export class CopilotController {
     return controller.signal;
   }
 
+  private handleError(err: any) {
+    if (err instanceof Error) {
+      const ret = {
+        message: err.message,
+        status: (err as any).status,
+      };
+      if (err instanceof HttpException) {
+        ret.status = err.getStatus();
+      }
+    }
+    return err;
+  }
+
   @Get('/chat/:sessionId')
   async chat(
     @CurrentUser() user: CurrentUser,
@@ -138,7 +154,14 @@ export class CopilotController {
     @Query('messageId') messageId: string,
     @Query() params: Record<string, string>
   ): Promise<Observable<ChatEvent>> {
-    await this.chatSession.checkQuota(user.id);
+    try {
+      await this.chatSession.checkQuota(user.id);
+    } catch (err) {
+      return of({
+        type: 'error' as const,
+        data: this.handleError(err),
+      });
+    }
 
     const model = await this.chatSession.get(sessionId).then(s => s?.model);
     const provider = this.provider.getProviderByCapability(
@@ -178,6 +201,12 @@ export class CopilotController {
             switchMap(() => EMPTY)
           )
         )
+      ),
+      catchError(err =>
+        of({
+          type: 'error' as const,
+          data: this.handleError(err),
+        })
       )
     );
   }
@@ -190,7 +219,14 @@ export class CopilotController {
     @Query('messageId') messageId: string,
     @Query() params: Record<string, string>
   ): Promise<Observable<ChatEvent>> {
-    await this.chatSession.checkQuota(user.id);
+    try {
+      await this.chatSession.checkQuota(user.id);
+    } catch (err) {
+      return of({
+        type: 'error' as const,
+        data: this.handleError(err),
+      });
+    }
 
     const hasAttachment = await this.hasAttachment(sessionId, messageId);
     const model = await this.chatSession.get(sessionId).then(s => s?.model);
@@ -245,6 +281,12 @@ export class CopilotController {
             switchMap(() => EMPTY)
           )
         )
+      ),
+      catchError(err =>
+        of({
+          type: 'error' as const,
+          data: this.handleError(err),
+        })
       )
     );
   }