IMAP Fixes & Improvements (#13582)

- Removed caching of IMAP client as it was problematic and didn't really
matter because of our pipeline, also added retry logic to client.
- Some code refactoring and cleanup
- Breaking changes:
   - Removed messageChannel.syncCursor
   - Uses messageFolder.syncCursor instead
- Removed Date based sync cursor in favor of message UID which is much
more reliable and precise

---------

Co-authored-by: Charles Bochet <charles@twenty.com>

packages/twenty-front/src/modules/settings/accounts/hooks/useImapSmtpCaldavConnectionForm.ts

@@ -20,6 +20,7 @@ import {
   connectionImapSmtpCalDav,
   isProtocolConfigured,
 } from '@/settings/accounts/validation-schemas/connectionImapSmtpCalDav';
+import { ApolloError } from '@apollo/client';
 import { isDefined } from 'twenty-shared/utils';
 import {
   ConnectedImapSmtpCaldavAccount,
@@ -149,12 +150,9 @@ export const useImapSmtpCaldavConnectionForm = ({
         enqueueSuccessSnackBar({ message: successMessage });
         navigate(SettingsPath.Accounts);
       } catch (error) {
-        const errorMessage =
-          error instanceof Error
-            ? error.message
-            : 'An unexpected error occurred';
-
-        enqueueErrorSnackBar({ message: errorMessage });
+        enqueueErrorSnackBar({
+          apolloError: error instanceof ApolloError ? error : undefined,
+        });
       }
     },
     [

packages/twenty-server/src/engine/twenty-orm/repository/workspace-insert-query-builder.ts

@@ -175,7 +175,11 @@ export class WorkspaceInsertQueryBuilder<
       const resultWithoutInsertionExtraColumns = result.raw.map(
         (rawResult: Record<string, string>) =>
           Object.keys(rawResult)
-            .filter((key) => this.expressionMap.returning.includes(key))
+            .filter(
+              (key) =>
+                this.expressionMap.returning.includes(key) ||
+                this.expressionMap.returning === '*',
+            )
             .reduce((filtered: Record<string, string>, key) => {
               filtered[key] = rawResult[key];
 

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/providers/imap-client.provider.ts

@@ -7,34 +7,67 @@ import { isDefined } from 'twenty-shared/utils';
 import { ImapSmtpCaldavParams } from 'src/engine/core-modules/imap-smtp-caldav-connection/types/imap-smtp-caldav-connection.type';
 import { ConnectedAccountWorkspaceEntity } from 'src/modules/connected-account/standard-objects/connected-account.workspace-entity';
 
-interface ImapClientInstance {
-  client: ImapFlow;
-  isReady: boolean;
-}
+type ConnectedAccountIdentifier = Pick<
+  ConnectedAccountWorkspaceEntity,
+  'id' | 'provider' | 'connectionParameters' | 'handle'
+>;
 
 @Injectable()
 export class ImapClientProvider {
   private readonly logger = new Logger(ImapClientProvider.name);
-  private readonly clientInstances = new Map<string, ImapClientInstance>();
+
+  private static readonly RETRY_ATTEMPTS = 3;
+  private static readonly RETRY_DELAY_MS = 1000;
+  private static readonly CONNECTION_TIMEOUT_MS = 30000;
 
   constructor() {}
 
   async getClient(
-    connectedAccount: Pick<
-      ConnectedAccountWorkspaceEntity,
-      'id' | 'provider' | 'connectionParameters' | 'handle'
-    >,
+    connectedAccount: ConnectedAccountIdentifier,
   ): Promise<ImapFlow> {
-    const cacheKey = `${connectedAccount.id}`;
+    return this.createConnectionWithRetry(connectedAccount);
+  }
 
-    if (this.clientInstances.has(cacheKey)) {
-      const instance = this.clientInstances.get(cacheKey);
-
-      if (instance?.isReady) {
-        return instance.client;
-      }
+  async closeClient(client: ImapFlow): Promise<void> {
+    try {
+      await client.logout();
+      this.logger.log('Closed IMAP client');
+    } catch (error) {
+      this.logger.error(`Error closing IMAP client: ${error.message}`);
     }
+  }
 
+  private async createConnectionWithRetry(
+    connectedAccount: ConnectedAccountIdentifier,
+    attempt = 1,
+  ): Promise<ImapFlow> {
+    try {
+      return await this.createConnection(connectedAccount);
+    } catch (error) {
+      if (attempt < ImapClientProvider.RETRY_ATTEMPTS) {
+        const delay = ImapClientProvider.RETRY_DELAY_MS * attempt;
+
+        this.logger.warn(
+          `IMAP connection attempt ${attempt} failed for ${connectedAccount.handle}, retrying in ${delay}ms: ${error.message}`,
+        );
+
+        await this.delay(delay);
+
+        return this.createConnectionWithRetry(connectedAccount, attempt + 1);
+      }
+
+      this.logger.error(
+        `Failed to establish IMAP connection for ${connectedAccount.handle} after ${ImapClientProvider.RETRY_ATTEMPTS} attempts: ${error.message}`,
+        error.stack,
+      );
+
+      throw error;
+    }
+  }
+
+  private async createConnection(
+    connectedAccount: ConnectedAccountIdentifier,
+  ): Promise<ImapFlow> {
     if (
       connectedAccount.provider !== ConnectedAccountProvider.IMAP_SMTP_CALDAV ||
       !isDefined(connectedAccount.connectionParameters?.IMAP)
@@ -46,22 +79,38 @@ export class ImapClientProvider {
       (connectedAccount.connectionParameters as unknown as ImapSmtpCaldavParams) ||
       {};
 
-    const client = new ImapFlow({
-      host: connectionParameters.IMAP?.host || '',
-      port: connectionParameters.IMAP?.port || 993,
-      secure: connectionParameters.IMAP?.secure,
-      auth: {
-        user: connectedAccount.handle,
-        pass: connectionParameters.IMAP?.password || '',
-      },
-      logger: false,
-      tls: {
-        rejectUnauthorized: false,
-      },
-    });
+    let client: ImapFlow | null = null;
+    let timeoutId: NodeJS.Timeout | null = null;
 
     try {
-      await client.connect();
+      client = new ImapFlow({
+        host: connectionParameters.IMAP?.host || '',
+        port: connectionParameters.IMAP?.port || 993,
+        secure: connectionParameters.IMAP?.secure,
+        auth: {
+          user: connectedAccount.handle,
+          pass: connectionParameters.IMAP?.password || '',
+        },
+        logger: false,
+        tls: {
+          rejectUnauthorized: false,
+        },
+      });
+
+      const connectionPromise = client.connect();
+      const timeoutPromise = new Promise<never>((_, reject) => {
+        timeoutId = setTimeout(
+          () => reject(new Error('Connection timeout')),
+          ImapClientProvider.CONNECTION_TIMEOUT_MS,
+        );
+      });
+
+      await Promise.race([connectionPromise, timeoutPromise]);
+
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+        timeoutId = null;
+      }
 
       this.logger.log(
         `Connected to IMAP server for ${connectedAccount.handle}`,
@@ -77,34 +126,27 @@ export class ImapClientProvider {
         this.logger.warn(`Failed to list mailboxes: ${error.message}`);
       }
 
-      this.clientInstances.set(cacheKey, {
-        client,
-        isReady: true,
-      });
-
       return client;
     } catch (error) {
-      this.logger.error(
-        `Failed to connect to IMAP server: ${error.message}`,
-        error.stack,
-      );
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+      }
+
+      if (client) {
+        try {
+          await client.logout();
+        } catch (cleanupError) {
+          this.logger.warn(
+            `Failed to cleanup client after connection error: ${cleanupError.message}`,
+          );
+        }
+      }
+
       throw error;
     }
   }
 
-  async closeClient(connectedAccountId: string): Promise<void> {
-    const cacheKey = `${connectedAccountId}`;
-    const instance = this.clientInstances.get(cacheKey);
-
-    if (instance?.isReady) {
-      try {
-        await instance.client.logout();
-        this.logger.log('Closed IMAP client');
-      } catch (error) {
-        this.logger.error(`Error closing IMAP client: ${error.message}`);
-      } finally {
-        this.clientInstances.delete(cacheKey);
-      }
-    }
+  private delay(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
   }
 }

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/services/imap-fetch-by-batch.service.ts

@@ -1,13 +1,8 @@
 import { Injectable, Logger } from '@nestjs/common';
 
-import { ImapFlow } from 'imapflow';
-
 import { ConnectedAccountWorkspaceEntity } from 'src/modules/connected-account/standard-objects/connected-account.workspace-entity';
 import { ImapClientProvider } from 'src/modules/messaging/message-import-manager/drivers/imap/providers/imap-client.provider';
-import {
-  ImapMessageLocatorService,
-  MessageLocation,
-} from 'src/modules/messaging/message-import-manager/drivers/imap/services/imap-message-locator.service';
+import { ImapMessageLocatorService } from 'src/modules/messaging/message-import-manager/drivers/imap/services/imap-message-locator.service';
 import {
   ImapMessageProcessorService,
   MessageFetchResult,
@@ -27,10 +22,6 @@ type FetchAllResult = {
 export class ImapFetchByBatchService {
   private readonly logger = new Logger(ImapFetchByBatchService.name);
 
-  private static readonly RETRY_ATTEMPTS = 2;
-  private static readonly RETRY_DELAY_MS = 1000;
-  private static readonly BATCH_LIMIT = 20;
-
   constructor(
     private readonly imapClientProvider: ImapClientProvider,
     private readonly imapMessageLocatorService: ImapMessageLocatorService,
@@ -41,6 +32,7 @@ export class ImapFetchByBatchService {
     messageIds: string[],
     connectedAccount: ConnectedAccount,
   ): Promise<FetchAllResult> {
+    const batchLimit = 20;
     const batchResults: MessageFetchResult[][] = [];
     const messageIdsByBatch: string[][] = [];
 
@@ -48,100 +40,56 @@ export class ImapFetchByBatchService {
       `Starting optimized batch fetch for ${messageIds.length} messages`,
     );
 
-    let client: ImapFlow | null = null;
+    const client = await this.imapClientProvider.getClient(connectedAccount);
 
     try {
-      client = await this.imapClientProvider.getClient(connectedAccount);
-
       const messageLocations =
         await this.imapMessageLocatorService.locateAllMessages(
           messageIds,
           client,
         );
 
-      const batches = this.chunkArray(
-        messageIds,
-        ImapFetchByBatchService.BATCH_LIMIT,
-      );
+      for (let i = 0; i < messageIds.length; i += batchLimit) {
+        const batchMessageIds = messageIds.slice(i, i + batchLimit);
 
-      let processedCount = 0;
+        messageIdsByBatch.push(batchMessageIds);
 
-      for (const batch of batches) {
-        const batchResult = await this.fetchBatchWithRetry(
-          batch,
-          messageLocations,
-          client,
-        );
+        try {
+          const batchResult =
+            await this.imapMessageProcessorService.processMessagesByIds(
+              batchMessageIds,
+              messageLocations,
+              client,
+            );
 
-        batchResults.push(batchResult);
-        messageIdsByBatch.push(batch);
+          batchResults.push(batchResult);
 
-        processedCount += batch.length;
-        this.logger.log(
-          `Fetched ${processedCount}/${messageIds.length} messages`,
-        );
+          this.logger.log(
+            `Fetched batch ${Math.floor(i / batchLimit) + 1}/${Math.ceil(messageIds.length / batchLimit)} (${batchMessageIds.length} messages)`,
+          );
+        } catch (error) {
+          this.logger.error(
+            `Batch fetch failed for batch starting at index ${i}: ${error.message}`,
+          );
+
+          const errorResults =
+            this.imapMessageProcessorService.createErrorResults(
+              batchMessageIds,
+              error as Error,
+            );
+
+          batchResults.push(errorResults);
+        }
       }
 
-      return { messageIdsByBatch, batchResults };
+      return {
+        messageIdsByBatch,
+        batchResults,
+      };
     } finally {
       if (client) {
-        await this.imapClientProvider.closeClient(connectedAccount.id);
+        await this.imapClientProvider.closeClient(client);
       }
     }
   }
-
-  private async fetchBatchWithRetry(
-    messageIds: string[],
-    messageLocations: Map<string, MessageLocation>,
-    client: ImapFlow,
-    attempt = 1,
-  ): Promise<MessageFetchResult[]> {
-    try {
-      return await this.imapMessageProcessorService.processMessagesByIds(
-        messageIds,
-        messageLocations,
-        client,
-      );
-    } catch (error) {
-      if (attempt < ImapFetchByBatchService.RETRY_ATTEMPTS) {
-        const delay = ImapFetchByBatchService.RETRY_DELAY_MS * attempt;
-
-        this.logger.warn(
-          `Batch fetch attempt ${attempt} failed, retrying in ${delay}ms: ${error.message}`,
-        );
-
-        await this.delay(delay);
-
-        return this.fetchBatchWithRetry(
-          messageIds,
-          messageLocations,
-          client,
-          attempt + 1,
-        );
-      }
-
-      this.logger.error(
-        `Batch fetch failed after ${ImapFetchByBatchService.RETRY_ATTEMPTS} attempts: ${error.message}`,
-      );
-
-      return this.imapMessageProcessorService.createErrorResults(
-        messageIds,
-        error as Error,
-      );
-    }
-  }
-
-  private chunkArray<T>(array: T[], chunkSize: number): T[][] {
-    const chunks: T[][] = [];
-
-    for (let i = 0; i < array.length; i += chunkSize) {
-      chunks.push(array.slice(i, i + chunkSize));
-    }
-
-    return chunks;
-  }
-
-  private delay(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-  }
 }

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/services/imap-get-message-list.service.ts

@@ -2,6 +2,7 @@ import { Injectable, Logger } from '@nestjs/common';
 
 import { ImapFlow } from 'imapflow';
 
+import { MessageFolderWorkspaceEntity } from 'src/modules/messaging/common/standard-objects/message-folder.workspace-entity';
 import { ImapClientProvider } from 'src/modules/messaging/message-import-manager/drivers/imap/providers/imap-client.provider';
 import { ImapHandleErrorService } from 'src/modules/messaging/message-import-manager/drivers/imap/services/imap-handle-error.service';
 import { MessageFolderName } from 'src/modules/messaging/message-import-manager/drivers/imap/types/folders';
@@ -21,13 +22,14 @@ export class ImapGetMessageListService {
     private readonly imapHandleErrorService: ImapHandleErrorService,
   ) {}
 
-  async getMessageLists({
-    messageChannel,
+  public async getMessageLists({
     connectedAccount,
     messageFolders,
   }: GetMessageListsArgs): Promise<GetMessageListsResponse> {
+    let client: ImapFlow | null = null;
+
     try {
-      const client = await this.imapClientProvider.getClient(connectedAccount);
+      client = await this.imapClientProvider.getClient(connectedAccount);
       const result: GetMessageListsResponse = [];
 
       for (const folder of messageFolders) {
@@ -38,10 +40,10 @@ export class ImapGetMessageListService {
         }
 
         try {
-          const response = await this.getMessageListForMailbox(
+          const response = await this.getMessageList(
             client,
             mailboxName,
-            folder.syncCursor,
+            folder,
           );
 
           result.push({
@@ -52,6 +54,14 @@ export class ImapGetMessageListService {
           this.logger.warn(
             `Error fetching from folder ${folder.name} (${mailboxName}): ${error.message}. Continuing with other folders.`,
           );
+
+          result.push({
+            messageExternalIds: [],
+            nextSyncCursor: folder.syncCursor || '',
+            previousSyncCursor: folder.syncCursor,
+            messageExternalIdsToDelete: [],
+            folderId: folder.id,
+          });
         }
       }
 
@@ -64,20 +74,47 @@ export class ImapGetMessageListService {
 
       this.imapHandleErrorService.handleImapMessageListFetchError(error);
 
-      return [
-        {
-          messageExternalIds: [],
-          nextSyncCursor: messageChannel.syncCursor || '',
-          previousSyncCursor: messageChannel.syncCursor,
-          messageExternalIdsToDelete: [],
-          folderId: undefined,
-        },
-      ];
+      return messageFolders.map((folder) => ({
+        messageExternalIds: [],
+        nextSyncCursor: folder.syncCursor || '',
+        previousSyncCursor: folder.syncCursor,
+        messageExternalIdsToDelete: [],
+        folderId: folder.id,
+      }));
     } finally {
-      await this.imapClientProvider.closeClient(connectedAccount.id);
+      if (client) {
+        await this.imapClientProvider.closeClient(client);
+      }
     }
   }
 
+  public async getMessageList(
+    client: ImapFlow,
+    mailbox: string,
+    messageFolder: Pick<MessageFolderWorkspaceEntity, 'syncCursor'>,
+  ): Promise<GetOneMessageListResponse> {
+    const messages = await this.getMessagesFromMailbox(
+      client,
+      mailbox,
+      messageFolder.syncCursor,
+    );
+
+    messages.sort((a, b) => parseInt(b.uid) - parseInt(a.uid));
+
+    const messageExternalIds = messages.map((message) => message.id);
+
+    const nextSyncCursor =
+      messages.length > 0 ? messages[0].uid : messageFolder.syncCursor || '';
+
+    return {
+      messageExternalIds,
+      nextSyncCursor,
+      previousSyncCursor: messageFolder.syncCursor || '',
+      messageExternalIdsToDelete: [],
+      folderId: undefined,
+    };
+  }
+
   private async getMailboxName(
     client: ImapFlow,
     folderName: string,
@@ -101,36 +138,11 @@ export class ImapGetMessageListService {
     return folderName;
   }
 
-  private async getMessageListForMailbox(
-    client: ImapFlow,
-    mailbox: string,
-    cursor?: string,
-  ): Promise<GetOneMessageListResponse> {
-    const messages = await this.getMessagesFromMailbox(client, mailbox, cursor);
-
-    messages.sort(
-      (a, b) => new Date(b.date).getTime() - new Date(a.date).getTime(),
-    );
-
-    const messageExternalIds = messages.map((message) => message.id);
-
-    const nextSyncCursor =
-      messages.length > 0 ? messages[messages.length - 1].date : cursor || '';
-
-    return {
-      messageExternalIds,
-      nextSyncCursor,
-      previousSyncCursor: cursor || '',
-      messageExternalIdsToDelete: [],
-      folderId: undefined,
-    };
-  }
-
   private async getMessagesFromMailbox(
     client: ImapFlow,
     mailbox: string,
     cursor?: string,
-  ): Promise<{ id: string; date: string }[]> {
+  ): Promise<{ id: string; uid: string }[]> {
     let lock;
 
     try {
@@ -139,27 +151,25 @@ export class ImapGetMessageListService {
       let searchOptions = {};
 
       if (cursor) {
-        searchOptions = {
-          since: new Date(cursor),
-        };
+        const cursorUid = parseInt(cursor);
+
+        if (!isNaN(cursorUid)) {
+          searchOptions = {
+            uid: `${cursorUid + 1}:*`,
+          };
+        }
       }
 
-      const messages: { id: string; date: string }[] = [];
+      const messages: { id: string; uid: string }[] = [];
 
       for await (const message of client.fetch(searchOptions, {
         envelope: true,
+        uid: true,
       })) {
-        if (message.envelope?.messageId) {
-          const messageDate = message.envelope.date
-            ? new Date(message.envelope.date)
-            : new Date();
-          const validDate = isNaN(messageDate.getTime())
-            ? new Date()
-            : messageDate;
-
+        if (message.envelope?.messageId && message.uid) {
           messages.push({
             id: message.envelope.messageId,
-            date: validDate.toISOString(),
+            uid: message.uid.toString(),
           });
         }
       }

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/services/imap-get-messages.service.ts

@@ -16,6 +16,11 @@ import { sanitizeString } from 'src/modules/messaging/message-import-manager/uti
 
 type AddressType = 'from' | 'to' | 'cc' | 'bcc';
 
+type ConnectedAccountType = Pick<
+  ConnectedAccountWorkspaceEntity,
+  'id' | 'provider' | 'handle' | 'handleAliases' | 'connectionParameters'
+>;
+
 @Injectable()
 export class ImapGetMessagesService {
   private readonly logger = new Logger(ImapGetMessagesService.name);
@@ -24,36 +29,40 @@ export class ImapGetMessagesService {
 
   async getMessages(
     messageIds: string[],
-    connectedAccount: Pick<
-      ConnectedAccountWorkspaceEntity,
-      'id' | 'provider' | 'handle' | 'handleAliases' | 'connectionParameters'
-    >,
+    connectedAccount: ConnectedAccountType,
   ): Promise<MessageWithParticipants[]> {
     if (!messageIds.length) {
       return [];
     }
 
-    const { messageIdsByBatch, batchResults } =
-      await this.fetchByBatchService.fetchAllByBatches(
-        messageIds,
-        connectedAccount,
-      );
+    const { batchResults } = await this.fetchByBatchService.fetchAllByBatches(
+      messageIds,
+      connectedAccount,
+    );
 
     this.logger.log(`IMAP fetch completed`);
 
-    const messages = batchResults.flatMap((batchResult, index) => {
-      return this.formatBatchResultAsMessages(
-        messageIdsByBatch[index],
-        batchResult,
-        connectedAccount,
-      );
-    });
+    const messages = this.formatBatchResponsesAsMessages(
+      batchResults,
+      connectedAccount,
+    );
 
     return messages;
   }
 
-  private formatBatchResultAsMessages(
-    messageIds: string[],
+  public formatBatchResponsesAsMessages(
+    batchResults: MessageFetchResult[][],
+    connectedAccount: Pick<
+      ConnectedAccountWorkspaceEntity,
+      'handle' | 'handleAliases'
+    >,
+  ): MessageWithParticipants[] {
+    return batchResults.flatMap((batchResult) => {
+      return this.formatBatchResponseAsMessages(batchResult, connectedAccount);
+    });
+  }
+
+  private formatBatchResponseAsMessages(
     batchResults: MessageFetchResult[],
     connectedAccount: Pick<
       ConnectedAccountWorkspaceEntity,

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/utils/parse-imap-error.util.ts

@@ -2,6 +2,7 @@ import {
   MessageImportDriverException,
   MessageImportDriverExceptionCode,
 } from 'src/modules/messaging/message-import-manager/drivers/exceptions/message-import-driver.exception';
+import { MessageNetworkExceptionCode } from 'src/modules/messaging/message-import-manager/drivers/exceptions/message-network.exception';
 import { isImapFlowError } from 'src/modules/messaging/message-import-manager/drivers/imap/utils/is-imap-flow-error.util';
 
 export const parseImapError = (
@@ -15,6 +16,20 @@ export const parseImapError = (
     return null;
   }
 
+  if (error.message.includes('Connection not available')) {
+    return new MessageImportDriverException(
+      `IMAP client not available: ${error.message}`,
+      MessageImportDriverExceptionCode.CLIENT_NOT_AVAILABLE,
+    );
+  }
+
+  if (error.message.includes('timeout')) {
+    return new MessageImportDriverException(
+      `IMAP connection timeout: ${error.message}`,
+      MessageNetworkExceptionCode.ETIMEDOUT,
+    );
+  }
+
   if (error.code === 'ECONNREFUSED' || error.message === 'Failed to connect') {
     return new MessageImportDriverException(
       `IMAP connection error: ${error.message}`,

packages/twenty-server/src/modules/messaging/message-import-manager/drivers/imap/utils/parse-imap-message-list-fetch-error.util.ts

@@ -57,7 +57,7 @@ export const parseImapMessageListFetchError = (
   }
 
   return new MessageImportDriverException(
-    `Unknown IMAP message list fetch error: ${errorMessage}`,
+    `Unknown IMAP message list fetch error: code: ${error.code} | responseText: ${error.responseText} | executedCommand: ${error.executedCommand}`,
     MessageImportDriverExceptionCode.UNKNOWN,
   );
 };