add auth retry for data chatbot

Author: ArtemHoruzhenko

redisinsight/api/src/modules/ai/query/ai-query.service.ts

@@ -121,7 +121,7 @@ export class AiQueryService {
         accountId: auth.accountId,
       });
 
-      socket = await this.aiQueryProvider.getSocket(auth);
+      socket = await this.aiQueryProvider.getSocket(sessionMetadata, auth);
 
       socket.on(AiQueryWsEvents.REPLY_CHUNK, (chunk) => {
         answer.content += chunk;

redisinsight/api/src/modules/ai/query/exceptions/ai-query.error.handler.ts

@@ -1,4 +1,5 @@
 import { AxiosError } from 'axios';
+import { get } from 'lodash';
 import { HttpException } from '@nestjs/common';
 import {
   AiQueryUnauthorizedException,
@@ -13,11 +14,12 @@ export const wrapAiQueryError = (error: AxiosError, message?: string): HttpExcep
     return error;
   }
 
-  const { response } = error;
+  // TransportError or Axios error
+  const response = get(error, ['description', 'target', '_req', 'res'], error.response);
 
   if (response) {
     const errorOptions = { cause: new Error(response?.data as string) };
-    switch (response?.status) {
+    switch (response?.status || response?.statusCode) {
       case 401:
         return new AiQueryUnauthorizedException(message, errorOptions);
       case 403:

redisinsight/api/src/modules/ai/query/providers/ai-query.provider.ts

@@ -3,34 +3,47 @@ import config, { Config } from 'src/utils/config';
 import { Injectable, Logger } from '@nestjs/common';
 import { AiQueryAuthData } from 'src/modules/ai/query/models/ai-query.auth-data';
 import { AiQueryWsEvents } from 'src/modules/ai/query/models';
+import { AiQueryAuthProvider } from 'src/modules/ai/query/providers/auth/ai-query-auth.provider';
+import { SessionMetadata } from 'src/common/models';
+import { wrapAiQueryError } from 'src/modules/ai/query/exceptions';
 
 const aiConfig = config.get('ai') as Config['ai'];
 
 @Injectable()
 export class AiQueryProvider {
   private readonly logger = new Logger('AiQueryProvider');
 
-  async getSocket(auth: AiQueryAuthData): Promise<Socket> {
-    return new Promise((resolve, reject) => {
-      const socket = io(aiConfig.querySocketUrl, {
-        path: aiConfig.querySocketPath,
-        reconnection: false,
-        transports: ['websocket'],
-        extraHeaders: {
-          'X-Csrf-Token': auth.csrf,
-          Cookie: `JSESSIONID=${auth.sessionId}`,
-        },
-      });
+  constructor(
+    private readonly aiQueryAuthProvider: AiQueryAuthProvider,
+  ) {}
 
-      socket.on(AiQueryWsEvents.CONNECT_ERROR, (e) => {
-        this.logger.error('Unable to establish AI socket connection', e);
-        reject(e);
-      });
+  async getSocket(sessionMetadata: SessionMetadata, auth: AiQueryAuthData): Promise<Socket> {
+    return this.aiQueryAuthProvider.callWithAuthRetry(sessionMetadata, async () => {
+      try {
+        return await new Promise((resolve, reject) => {
+          const socket = io(aiConfig.querySocketUrl, {
+            path: aiConfig.querySocketPath,
+            reconnection: false,
+            transports: ['websocket'],
+            extraHeaders: {
+              'X-Csrf-Token': auth.csrf,
+              Cookie: `JSESSIONID=${auth.sessionId}`,
+            },
+          });
 
-      socket.on(AiQueryWsEvents.CONNECT, async () => {
-        this.logger.debug('AI socket connection established');
-        resolve(socket);
-      });
+          socket.on(AiQueryWsEvents.CONNECT_ERROR, (e) => {
+            this.logger.error('Unable to establish AI socket connection', e);
+            reject(e);
+          });
+
+          socket.on(AiQueryWsEvents.CONNECT, async () => {
+            this.logger.debug('AI socket connection established');
+            resolve(socket);
+          });
+        });
+      } catch (e) {
+        throw wrapAiQueryError(e, 'Unable to establish connection');
+      }
     });
   }
 }

redisinsight/api/src/modules/ai/query/providers/auth/ai-query-auth.provider.ts

@@ -3,4 +3,5 @@ import { AiQueryAuthData } from 'src/modules/ai/query/models/ai-query.auth-data'
 
 export abstract class AiQueryAuthProvider {
   abstract getAuthData(sessionMetadata: SessionMetadata): Promise<AiQueryAuthData>;
+  abstract callWithAuthRetry(sessionId: SessionMetadata, fn: () => Promise<any>, retries?: number): Promise<any>;
 }

redisinsight/api/src/modules/ai/query/providers/auth/local.ai-query-auth.provider.ts

@@ -3,6 +3,8 @@ import { SessionMetadata } from 'src/common/models';
 import { AiQueryAuthData } from 'src/modules/ai/query/models/ai-query.auth-data';
 import { AiQueryAuthProvider } from 'src/modules/ai/query/providers/auth/ai-query-auth.provider';
 import { CloudUserApiService } from 'src/modules/cloud/user/cloud-user.api.service';
+import { CloudApiUnauthorizedException } from 'src/modules/cloud/common/exceptions';
+import { AiQueryUnauthorizedException } from 'src/modules/ai/query/exceptions';
 
 @Injectable()
 export class LocalAiQueryAuthProvider extends AiQueryAuthProvider {
@@ -22,4 +24,17 @@ export class LocalAiQueryAuthProvider extends AiQueryAuthProvider {
       accountId: `${session.user.currentAccountId}`,
     };
   }
+
+  async callWithAuthRetry(sessionMetadata: SessionMetadata, fn: () => Promise<any>, retries = 1) {
+    try {
+      return await fn();
+    } catch (e) {
+      if (retries > 0 && (e instanceof CloudApiUnauthorizedException || e instanceof AiQueryUnauthorizedException)) {
+        await this.cloudUserApiService.invalidateApiSession(sessionMetadata).catch(() => {});
+        return this.callWithAuthRetry(sessionMetadata, fn, retries - 1);
+      }
+
+      throw e;
+    }
+  }
 }

redisinsight/api/src/modules/cloud/user/cloud-user.api.service.ts

@@ -213,6 +213,16 @@ export class CloudUserApiService {
     });
   }
 
+  /**
+   * Invalidate user SM API session
+   * @param sessionMetadata
+   */
+  async invalidateApiSession(
+    sessionMetadata: SessionMetadata,
+  ): Promise<void> {
+    await this.sessionService.invalidateApiSession(sessionMetadata.sessionId);
+  }
+
   /**
    * Select current account to work with
    * @param sessionMetadata