feat: allow kill switch to disable saving agent configs and dsviews (#…

…9356) * feat: allow kill switch to disable saving agent configs and dsviews * block at endpoint level * fix * add assistant builder tooltip * ds views tooltips --------- Co-authored-by: Henry Fontanier <[email protected]>
dust-tt · Dec 13, 2024 · 7e21c12 · 7e21c12
1 parent cafe8c1
commit 7e21c12
Show file tree

Hide file tree

Showing 13 changed files with 956 additions and 711 deletions.
diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -65,6 +65,7 @@ import {
   AppLayoutSimpleSaveCancelTitle,
 } from "@app/components/sparkle/AppLayoutTitle";
 import { isUpgraded } from "@app/lib/plans/plan_codes";
+import { useKillSwitches } from "@app/lib/swr/kill";
 
 function isValidTab(tab: string): tab is BuilderScreen {
   return BUILDER_SCREENS.includes(tab as BuilderScreen);
@@ -85,6 +86,9 @@ export default function AssistantBuilder({
   const router = useRouter();
   const sendNotification = useSendNotification();
 
+  const { killSwitches } = useKillSwitches();
+  const isSavingDisabled = killSwitches?.includes("save_agent_configurations");
+
   const defaultScope =
     flow === "workspace_assistants" ? "workspace" : "private";
   const [currentTab, setCurrentTab] = useHashParam(
@@ -380,8 +384,13 @@ export default function AssistantBuilder({
               onCancel={async () => {
                 await appLayoutBack(owner, router);
               }}
-              onSave={onAssistantSave}
+              onSave={isSavingDisabled ? undefined : onAssistantSave}
               isSaving={isSavingOrDeleting}
+              saveTooltip={
+                isSavingDisabled
+                  ? "Saving assistants is temporarily disabled and will be re-enabled shortly."
+                  : undefined
+              }
             />
           )
         }

diff --git a/front/components/spaces/EditSpaceManagedDatasourcesViews.tsx b/front/components/spaces/EditSpaceManagedDatasourcesViews.tsx
@@ -4,6 +4,7 @@ import {
   Dialog,
   InformationCircleIcon,
   PlusIcon,
+  Tooltip,
 } from "@dust-tt/sparkle";
 import { useSendNotification } from "@dust-tt/sparkle";
 import type {
@@ -23,6 +24,7 @@ import { RequestDataSourceModal } from "@app/components/data_source/RequestDataS
 import SpaceManagedDatasourcesViewsModal from "@app/components/spaces/SpaceManagedDatasourcesViewsModal";
 import { useAwaitableDialog } from "@app/hooks/useAwaitableDialog";
 import { getDisplayNameForDataSource, isManaged } from "@app/lib/data_sources";
+import { useKillSwitches } from "@app/lib/swr/kill";
 import {
   useSpaceDataSourceViews,
   useSpaceDataSourceViewsWithDetails,
@@ -54,6 +56,10 @@ export function EditSpaceManagedDataSourcesViews({
 
   const { AwaitableDialog, showDialog } = useAwaitableDialog();
 
+  const { killSwitches } = useKillSwitches();
+
+  const isSavingDisabled = killSwitches?.includes("save_data_source_views");
+
   // DataSources Views of the current space.
   const {
     spaceDataSourceViews,
@@ -264,6 +270,28 @@ export function EditSpaceManagedDataSourcesViews({
   if (isSystemSpaceDataSourceViewsLoading || isSpaceDataSourceViewsLoading) {
     return false;
   }
+
+  const addToSpaceButton = (
+    <Button
+      label={
+        dataSourceView
+          ? `Add data from ${getDisplayNameForDataSource(dataSourceView.dataSource)}`
+          : "Add data from connections"
+      }
+      variant="primary"
+      icon={PlusIcon}
+      size="sm"
+      onClick={() => {
+        if (systemSpaceDataSourceViews.length === 0) {
+          setShowNoConnectionDialog(true);
+        } else {
+          setShowDataSourcesModal(true);
+        }
+      }}
+      disabled={isSavingDisabled}
+    />
+  );
+
   return isAdmin ? (
     <>
       <SpaceManagedDatasourcesViewsModal
@@ -307,23 +335,14 @@ export function EditSpaceManagedDataSourcesViews({
         <p>You have no connection set up.</p>
       </Dialog>
       <AwaitableDialog />
-      <Button
-        label={
-          dataSourceView
-            ? `Add data from ${getDisplayNameForDataSource(dataSourceView.dataSource)}`
-            : "Add data from connections"
-        }
-        variant="primary"
-        icon={PlusIcon}
-        size="sm"
-        onClick={() => {
-          if (systemSpaceDataSourceViews.length === 0) {
-            setShowNoConnectionDialog(true);
-          } else {
-            setShowDataSourcesModal(true);
-          }
-        }}
-      />
+      {isSavingDisabled ? (
+        <Tooltip
+          trigger={addToSpaceButton}
+          label="Editing spaces is temporarily disabled and will be re-enabled shortly."
+        />
+      ) : (
+        addToSpaceButton
+      )}
     </>
   ) : (
     <RequestDataSourceModal

diff --git a/front/components/spaces/EditSpaceStaticDatasourcesViews.tsx b/front/components/spaces/EditSpaceStaticDatasourcesViews.tsx
@@ -11,6 +11,7 @@ import { useCallback, useState } from "react";
 
 import SpaceFolderModal from "@app/components/spaces/SpaceFolderModal";
 import SpaceWebsiteModal from "@app/components/spaces/SpaceWebsiteModal";
+import { useKillSwitches } from "@app/lib/swr/kill";
 
 interface EditSpaceStaticDatasourcesViewsProps {
   canWriteInSpace: boolean;
@@ -40,7 +41,9 @@ export function EditSpaceStaticDatasourcesViews({
   const router = useRouter();
   const [showDatasourceLimitPopup, setShowDatasourceLimitPopup] =
     useState(false);
+  const { killSwitches } = useKillSwitches();
 
+  const isSavingDisabled = killSwitches?.includes("save_data_source_views");
   const planDataSourcesLimit = plan.limits.dataSources.count;
 
   const checkLimitsAndOpenModal = useCallback(() => {
@@ -54,6 +57,15 @@ export function EditSpaceStaticDatasourcesViews({
     }
   }, [dataSources.length, planDataSourcesLimit, onOpen]);
 
+  const addToSpaceButton = (
+    <Button
+      label={`Add ${category}`}
+      onClick={checkLimitsAndOpenModal}
+      icon={PlusIcon}
+      disabled={!canWriteInSpace || isSavingDisabled}
+    />
+  );
+
   return (
     <>
       <Popup
@@ -89,12 +101,15 @@ export function EditSpaceStaticDatasourcesViews({
         />
       ) : null}
       {canWriteInSpace ? (
-        <Button
-          label={`Add ${category}`}
-          onClick={checkLimitsAndOpenModal}
-          icon={PlusIcon}
-          disabled={!canWriteInSpace}
-        />
+        isSavingDisabled ? (
+          <Tooltip
+            label="Editing spaces is temporarily disabled and will be re-enabled shortly."
+            side="top"
+            trigger={addToSpaceButton}
+          />
+        ) : (
+          addToSpaceButton
+        )
       ) : (
         <Tooltip
           label={
@@ -103,14 +118,7 @@ export function EditSpaceStaticDatasourcesViews({
               : `Only members of the space can add a ${category}.`
           }
           side="top"
-          trigger={
-            <Button
-              label={`Add ${category}`}
-              onClick={checkLimitsAndOpenModal}
-              icon={PlusIcon}
-              disabled={!canWriteInSpace}
-            />
-          }
+          trigger={addToSpaceButton}
         />
       )}
     </>

diff --git a/front/components/sparkle/AppLayoutTitle.tsx b/front/components/sparkle/AppLayoutTitle.tsx
@@ -22,11 +22,13 @@ export function AppLayoutSimpleSaveCancelTitle({
   onSave,
   onCancel,
   isSaving,
+  saveTooltip,
 }: {
   title: string;
   onSave?: () => void;
   onCancel: () => void;
   isSaving?: boolean;
+  saveTooltip?: string;
 }) {
   return (
     <BarHeader
@@ -37,6 +39,7 @@ export function AppLayoutSimpleSaveCancelTitle({
           onCancel={onCancel}
           onSave={onSave}
           isSaving={isSaving}
+          saveTooltip={saveTooltip}
         />
       }
       className="ml-10 lg:ml-0"

diff --git a/front/lib/resources/kill_switch_resource.ts b/front/lib/resources/kill_switch_resource.ts
@@ -45,11 +45,9 @@ export class KillSwitchResource extends BaseResource<KillSwitchModel> {
     });
   }
 
-  static async list(): Promise<KillSwitchResource[]> {
+  static async listEnabledKillSwitches(): Promise<KillSwitchType[]> {
     const killSwitches = await KillSwitchModel.findAll();
-    return killSwitches.map(
-      (ks) => new KillSwitchResource(KillSwitchModel, ks.get())
-    );
+    return killSwitches.map((ks) => ks.type);
   }
 
   async delete(): Promise<Result<number | undefined, Error>> {

diff --git a/front/lib/swr/kill.ts b/front/lib/swr/kill.ts
@@ -0,0 +1,22 @@
+// LABS - CAN BE REMOVED ANYTIME
+
+import type { Fetcher } from "swr";
+
+import { fetcher, useSWRWithDefaults } from "@app/lib/swr/swr";
+import type { GetKillSwitchesResponseBody } from "@app/pages/api/kill";
+
+export function useKillSwitches() {
+  const killSwitchesFetcher: Fetcher<GetKillSwitchesResponseBody> = fetcher;
+
+  const { data, error, mutate } = useSWRWithDefaults(
+    `/api/kill`,
+    killSwitchesFetcher
+  );
+
+  return {
+    killSwitches: data ? data.killSwitches : null,
+    isKillSwitchesLoading: !error && !data,
+    isKillSwitchesError: error,
+    mutateKillSwitches: mutate,
+  };
+}