From 24f2b8e9c3974cd1fb597e05719e830ac28fe809 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <daphne.popin@gmail.com>
Date: Fri, 15 Nov 2024 00:22:13 +0100
Subject: [PATCH] Display specific error message if fail to retrieve
 permissions caused by rate limit

---
 connectors/src/api/get_connector_permissions.ts | 13 +++++++++++++
 connectors/src/connectors/slack/index.ts        | 14 ++++++++++++--
 .../src/connectors/slack/lib/slack_client.ts    |  5 +----
 front/components/ContentNodeTree.tsx            | 17 ++++++++++++++---
 front/lib/api/data_source_view.ts               |  8 ++++++++
 front/lib/swr/connectors.ts                     |  5 ++++-
 .../[dsId]/managed/permissions/index.ts         | 10 ++++++++++
 7 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/connectors/src/api/get_connector_permissions.ts b/connectors/src/api/get_connector_permissions.ts
index 71dc3f09b56fe..bb3598f1c7fb7 100644
--- a/connectors/src/api/get_connector_permissions.ts
+++ b/connectors/src/api/get_connector_permissions.ts
@@ -7,6 +7,7 @@ import type { Request, Response } from "express";
 
 import { getConnectorManager } from "@connectors/connectors";
 import { augmentContentNodesWithParentIds } from "@connectors/lib/api/content_nodes";
+import { ProviderWorkflowError } from "@connectors/lib/error";
 import logger from "@connectors/logger/logger";
 import { apiError, withLogging } from "@connectors/logger/withlogging";
 import { ConnectorResource } from "@connectors/resources/connector_resource";
@@ -85,6 +86,18 @@ const _getConnectorPermissions = async (
   });
 
   if (pRes.isErr()) {
+    if (
+      pRes.error instanceof ProviderWorkflowError &&
+      pRes.error.type === "rate_limit_error"
+    ) {
+      return apiError(req, res, {
+        status_code: 429,
+        api_error: {
+          type: "connector_rate_limit_error",
+          message: pRes.error.message,
+        },
+      });
+    }
     return apiError(req, res, {
       status_code: 500,
       api_error: {
diff --git a/connectors/src/connectors/slack/index.ts b/connectors/src/connectors/slack/index.ts
index 3729371c2a30e..c06b4e13fec86 100644
--- a/connectors/src/connectors/slack/index.ts
+++ b/connectors/src/connectors/slack/index.ts
@@ -21,7 +21,10 @@ import {
   getSlackClient,
 } from "@connectors/connectors/slack/lib/slack_client";
 import { launchSlackSyncWorkflow } from "@connectors/connectors/slack/temporal/client.js";
-import { ExternalOAuthTokenError } from "@connectors/lib/error";
+import {
+  ExternalOAuthTokenError,
+  ProviderWorkflowError,
+} from "@connectors/lib/error";
 import { SlackChannel } from "@connectors/lib/models/slack";
 import { terminateAllWorkflowsForConnectorId } from "@connectors/lib/temporal";
 import logger from "@connectors/logger/logger";
@@ -274,7 +277,7 @@ export class SlackConnectorManager extends BaseConnectorManager<SlackConfigurati
     parentInternalId: string | null;
     filterPermission: ConnectorPermission | null;
     viewType: ContentNodesViewType;
-  }): Promise<Result<ContentNode[], Error>> {
+  }): Promise<Result<ContentNode[], Error | ProviderWorkflowError>> {
     if (parentInternalId) {
       return new Err(
         new Error(
@@ -404,6 +407,13 @@ export class SlackConnectorManager extends BaseConnectorManager<SlackConfigurati
           new Error("Slack token invalid. Please re-authorize Slack.")
         );
       }
+      if (e instanceof ProviderWorkflowError && e.type === "rate_limit_error") {
+        logger.error(
+          { connectorId: this.connectorId, error: e },
+          "Slack rate limit when retrieving permissions."
+        );
+        return new Err(e);
+      }
       throw e;
     }
   }
diff --git a/connectors/src/connectors/slack/lib/slack_client.ts b/connectors/src/connectors/slack/lib/slack_client.ts
index 1fce7f4b59267..439037115e626 100644
--- a/connectors/src/connectors/slack/lib/slack_client.ts
+++ b/connectors/src/connectors/slack/lib/slack_client.ts
@@ -65,10 +65,7 @@ export async function getSlackClient(
           // If we get rate limited, we throw a known error.
           // Note: a previous version using slackError.code === ErrorCode.RateLimitedError failed
           // see PR #2689 for details
-          if (
-            e instanceof Error &&
-            e.message.startsWith("A rate limit was exceeded")
-          ) {
+          if (e instanceof Error && e.message.includes("rate limit")) {
             try {
               Context.current().heartbeat();
               await Context.current().sleep("1 minute");
diff --git a/front/components/ContentNodeTree.tsx b/front/components/ContentNodeTree.tsx
index 615ddfea2e013..72fedc1477855 100644
--- a/front/components/ContentNodeTree.tsx
+++ b/front/components/ContentNodeTree.tsx
@@ -8,7 +8,7 @@ import {
   Tooltip,
   Tree,
 } from "@dust-tt/sparkle";
-import type { BaseContentNode } from "@dust-tt/types";
+import type { APIError, BaseContentNode } from "@dust-tt/types";
 import type { ReactNode } from "react";
 import React, { useCallback, useContext, useState } from "react";
 
@@ -35,6 +35,7 @@ export type UseResourcesHook = (parentId: string | null) => {
   resources: BaseContentNode[];
   isResourcesLoading: boolean;
   isResourcesError: boolean;
+  resourcesError: APIError | null;
 };
 
 export type ContentNodeTreeItemStatus = {
@@ -112,7 +113,7 @@ function ContentNodeTreeChildren({
 
   const { useResourcesHook, emptyComponent } = useContentNodeTreeContext();
 
-  const { resources, isResourcesLoading, isResourcesError } =
+  const { resources, isResourcesLoading, isResourcesError, resourcesError } =
     useResourcesHook(parentId);
 
   const filteredNodes = resources.filter(
@@ -152,7 +153,17 @@ function ContentNodeTreeChildren({
   if (isResourcesError) {
     return (
       <div className="text-sm text-warning">
-        Failed to retrieve permissions likely due to a revoked authorization.
+        {resourcesError?.type === "rate_limit_error" ? (
+          <>
+            Failed to retrieve permissions, we are currently rate limited by the
+            provider API.
+          </>
+        ) : (
+          <>
+            Failed to retrieve permissions likely due to a revoked
+            authorization.
+          </>
+        )}
       </div>
     );
   }
diff --git a/front/lib/api/data_source_view.ts b/front/lib/api/data_source_view.ts
index 5f63b7c691e3a..3571a726dc4f9 100644
--- a/front/lib/api/data_source_view.ts
+++ b/front/lib/api/data_source_view.ts
@@ -109,6 +109,14 @@ async function getContentNodesForManagedDataSourceView(
     });
 
     if (connectorsRes.isErr()) {
+      if (connectorsRes.error.type === "connector_rate_limit_error") {
+        return new Err(
+          new Error(
+            connectorsRes.error.message ??
+              "An error occurred while fetching the resources' content nodes."
+          )
+        );
+      }
       return new Err(
         new Error(
           "An error occurred while fetching the resources' children content nodes."
diff --git a/front/lib/swr/connectors.ts b/front/lib/swr/connectors.ts
index 0e13ae266d962..b8b08e134ea5f 100644
--- a/front/lib/swr/connectors.ts
+++ b/front/lib/swr/connectors.ts
@@ -1,5 +1,6 @@
 import { useSendNotification } from "@dust-tt/sparkle";
 import type {
+  APIError,
   ConnectorPermission,
   ContentNodesViewType,
   DataSourceType,
@@ -21,7 +22,8 @@ import type { GetDataSourcePermissionsResponseBody } from "@app/pages/api/w/[wId
 interface UseConnectorPermissionsReturn<IncludeParents extends boolean> {
   resources: GetDataSourcePermissionsResponseBody<IncludeParents>["resources"];
   isResourcesLoading: boolean;
-  isResourcesError: any;
+  isResourcesError: boolean;
+  resourcesError: APIError | null;
 }
 
 export function useConnectorPermissions<IncludeParents extends boolean>({
@@ -77,6 +79,7 @@ export function useConnectorPermissions<IncludeParents extends boolean>({
     ),
     isResourcesLoading: !error && !data,
     isResourcesError: error,
+    resourcesError: error ? (error.error as APIError) : null,
   };
 }
 
diff --git a/front/pages/api/w/[wId]/data_sources/[dsId]/managed/permissions/index.ts b/front/pages/api/w/[wId]/data_sources/[dsId]/managed/permissions/index.ts
index 33f0ee51e75b3..e117819a8a45e 100644
--- a/front/pages/api/w/[wId]/data_sources/[dsId]/managed/permissions/index.ts
+++ b/front/pages/api/w/[wId]/data_sources/[dsId]/managed/permissions/index.ts
@@ -272,6 +272,16 @@ export async function getManagedDataSourcePermissionsHandler(
     includeParents,
   });
   if (permissionsRes.isErr()) {
+    if (permissionsRes.error.type === "connector_rate_limit_error") {
+      return apiError(req, res, {
+        status_code: 429,
+        api_error: {
+          type: "rate_limit_error",
+          message:
+            "Rate limit error while retrieving the data source permissions",
+        },
+      });
+    }
     return apiError(req, res, {
       status_code: 500,
       api_error: {