From 45d8422ada8a484f5a6e306aeb41722afff4423b Mon Sep 17 00:00:00 2001
From: Henry Fontanier <henry@dust.tt>
Date: Thu, 12 Dec 2024 17:04:37 +0100
Subject: [PATCH] feat: add kill switches in poke

---
 front/components/poke/PokeNavbar.tsx          |   1 +
 front/lib/poke/types.ts                       |   8 +
 front/lib/resources/kill_switch_resource.ts   |  64 ++++++++
 .../resources/storage/models/kill_switches.ts |  43 ++++++
 front/migrations/db/migration_129.sql         |   9 ++
 front/pages/api/poke/kill.ts                  |  91 +++++++++++
 front/pages/poke/kill.tsx                     | 144 ++++++++++++++++++
 front/poke/swr/kill.ts                        |  19 +++
 8 files changed, 379 insertions(+)
 create mode 100644 front/lib/poke/types.ts
 create mode 100644 front/lib/resources/kill_switch_resource.ts
 create mode 100644 front/lib/resources/storage/models/kill_switches.ts
 create mode 100644 front/migrations/db/migration_129.sql
 create mode 100644 front/pages/api/poke/kill.ts
 create mode 100644 front/pages/poke/kill.tsx
 create mode 100644 front/poke/swr/kill.ts

diff --git a/front/components/poke/PokeNavbar.tsx b/front/components/poke/PokeNavbar.tsx
index 2675712a7027..1c0e49dc7090 100644
--- a/front/components/poke/PokeNavbar.tsx
+++ b/front/components/poke/PokeNavbar.tsx
@@ -29,6 +29,7 @@ export const PokeNavbar: React.FC = () => (
         <Button href="/poke/plans" variant="ghost" label="Plans" />
         <Button href="/poke/templates" variant="ghost" label="Templates" />
         <Button href="/poke/plugins" variant="ghost" label="Plugins" />
+        <Button href="/poke/kill" variant="ghost" label="Kill Switches" />
       </div>
     </div>
     <PokeSearchCommand />
diff --git a/front/lib/poke/types.ts b/front/lib/poke/types.ts
new file mode 100644
index 000000000000..82aa20fd3a48
--- /dev/null
+++ b/front/lib/poke/types.ts
@@ -0,0 +1,8 @@
+export const KILL_SWITCH_TYPES = [
+  "save_agent_configurations",
+  "save_data_source_views",
+] as const;
+export type KillSwitchType = (typeof KILL_SWITCH_TYPES)[number];
+export function isKillSwitchType(type: string): type is KillSwitchType {
+  return KILL_SWITCH_TYPES.includes(type as KillSwitchType);
+}
diff --git a/front/lib/resources/kill_switch_resource.ts b/front/lib/resources/kill_switch_resource.ts
new file mode 100644
index 000000000000..be76a4016668
--- /dev/null
+++ b/front/lib/resources/kill_switch_resource.ts
@@ -0,0 +1,64 @@
+import type { Result } from "@dust-tt/types";
+import { Ok } from "@dust-tt/types";
+import type { Attributes, ModelStatic } from "sequelize";
+
+import type { KillSwitchType } from "@app/lib/poke/types";
+import { BaseResource } from "@app/lib/resources/base_resource";
+import { KillSwitchModel } from "@app/lib/resources/storage/models/kill_switches";
+import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types";
+
+// Attributes are marked as read-only to reflect the stateless nature of our Resource.
+// This design will be moved up to BaseResource once we transition away from Sequelize.
+// eslint-disable-next-line @typescript-eslint/no-empty-interface
+// eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
+export interface KillSwitchResource
+  extends ReadonlyAttributesType<KillSwitchModel> {}
+// eslint-disable-next-line @typescript-eslint/no-unsafe-declaration-merging
+export class KillSwitchResource extends BaseResource<KillSwitchModel> {
+  static model: ModelStatic<KillSwitchModel> = KillSwitchModel;
+
+  constructor(
+    model: ModelStatic<KillSwitchModel>,
+    blob: Attributes<KillSwitchModel>
+  ) {
+    super(KillSwitchModel, blob);
+  }
+
+  static async enableKillSwitch(
+    type: KillSwitchType
+  ): Promise<KillSwitchResource> {
+    const ks =
+      (await KillSwitchModel.findOne({
+        where: {
+          type,
+        },
+      })) ?? (await KillSwitchModel.create({ type }));
+
+    return new KillSwitchResource(KillSwitchModel, ks.get());
+  }
+
+  static async disableKillSwitch(type: KillSwitchType): Promise<void> {
+    await KillSwitchModel.destroy({
+      where: {
+        type,
+      },
+    });
+  }
+
+  static async list(): Promise<KillSwitchResource[]> {
+    const killSwitches = await KillSwitchModel.findAll();
+    return killSwitches.map(
+      (ks) => new KillSwitchResource(KillSwitchModel, ks.get())
+    );
+  }
+
+  async delete(): Promise<Result<number | undefined, Error>> {
+    await this.model.destroy({
+      where: {
+        id: this.id,
+      },
+    });
+
+    return new Ok(this.id);
+  }
+}
diff --git a/front/lib/resources/storage/models/kill_switches.ts b/front/lib/resources/storage/models/kill_switches.ts
new file mode 100644
index 000000000000..180908e201bc
--- /dev/null
+++ b/front/lib/resources/storage/models/kill_switches.ts
@@ -0,0 +1,43 @@
+import type { CreationOptional } from "sequelize";
+import { DataTypes } from "sequelize";
+
+import type { KillSwitchType } from "@app/lib/poke/types";
+import { frontSequelize } from "@app/lib/resources/storage";
+import { BaseModel } from "@app/lib/resources/storage/wrappers";
+
+export class KillSwitchModel extends BaseModel<KillSwitchModel> {
+  declare id: CreationOptional<number>;
+  declare createdAt: CreationOptional<Date>;
+  declare updatedAt: CreationOptional<Date>;
+
+  declare type: KillSwitchType;
+}
+KillSwitchModel.init(
+  {
+    id: {
+      type: DataTypes.INTEGER,
+      autoIncrement: true,
+      primaryKey: true,
+    },
+    createdAt: {
+      type: DataTypes.DATE,
+      allowNull: false,
+      defaultValue: DataTypes.NOW,
+    },
+    updatedAt: {
+      type: DataTypes.DATE,
+      allowNull: false,
+      defaultValue: DataTypes.NOW,
+    },
+    type: {
+      type: DataTypes.STRING,
+      allowNull: false,
+    },
+  },
+  {
+    modelName: "kill_switches",
+    tableName: "kill_switches",
+    sequelize: frontSequelize,
+    indexes: [{ unique: true, fields: ["type"] }],
+  }
+);
diff --git a/front/migrations/db/migration_129.sql b/front/migrations/db/migration_129.sql
new file mode 100644
index 000000000000..171a8fef3783
--- /dev/null
+++ b/front/migrations/db/migration_129.sql
@@ -0,0 +1,9 @@
+CREATE TABLE IF NOT EXISTS "kill_switches" (
+  "id" BIGSERIAL PRIMARY KEY,
+  "createdAt" TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+  "updatedAt" TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+  "type" TEXT NOT NULL,
+  "enabled" BOOLEAN NOT NULL DEFAULT FALSE
+);
+
+CREATE UNIQUE INDEX IF NOT EXISTS "kill_switches_type_idx" ON "kill_switches" ("type");
\ No newline at end of file
diff --git a/front/pages/api/poke/kill.ts b/front/pages/api/poke/kill.ts
new file mode 100644
index 000000000000..a0da47177a3a
--- /dev/null
+++ b/front/pages/api/poke/kill.ts
@@ -0,0 +1,91 @@
+import type { WithAPIErrorResponse } from "@dust-tt/types";
+import { isLeft } from "fp-ts/lib/Either";
+import * as t from "io-ts";
+import * as reporter from "io-ts-reporters";
+import type { NextApiRequest, NextApiResponse } from "next";
+
+import { withSessionAuthentication } from "@app/lib/api/auth_wrappers";
+import { Authenticator } from "@app/lib/auth";
+import type { SessionWithUser } from "@app/lib/iam/provider";
+import type { KillSwitchType } from "@app/lib/poke/types";
+import { isKillSwitchType } from "@app/lib/poke/types";
+import { KillSwitchResource } from "@app/lib/resources/kill_switch_resource";
+import { apiError } from "@app/logger/withlogging";
+
+export type GetKillSwitchesResponseBody = {
+  killSwitches: KillSwitchType[];
+};
+
+const KillSwitchTypeSchema = t.type({
+  enabled: t.boolean,
+  type: t.string,
+});
+
+async function handler(
+  req: NextApiRequest,
+  res: NextApiResponse<
+    WithAPIErrorResponse<GetKillSwitchesResponseBody | { success: true }>
+  >,
+  session: SessionWithUser
+): Promise<void> {
+  const auth = await Authenticator.fromSuperUserSession(session, null);
+
+  if (!auth.isDustSuperUser()) {
+    return apiError(req, res, {
+      status_code: 404,
+      api_error: {
+        type: "user_not_found",
+        message: "Could not find the user.",
+      },
+    });
+  }
+
+  switch (req.method) {
+    case "GET":
+      const killSwitches = await KillSwitchResource.list();
+      return res
+        .status(200)
+        .json({ killSwitches: killSwitches.map((ks) => ks.type) });
+    case "POST":
+      const payloadValidation = KillSwitchTypeSchema.decode(req.body);
+      if (isLeft(payloadValidation)) {
+        const pathError = reporter.formatValidationErrors(
+          payloadValidation.left
+        );
+        return apiError(req, res, {
+          status_code: 400,
+          api_error: {
+            type: "invalid_request_error",
+            message: `The request body is invalid: ${pathError}`,
+          },
+        });
+      }
+      const { enabled, type } = payloadValidation.right;
+      if (!isKillSwitchType(type)) {
+        return apiError(req, res, {
+          status_code: 400,
+          api_error: {
+            type: "invalid_request_error",
+            message: `The request body is invalid: ${type} is not a valid kill switch type`,
+          },
+        });
+      }
+      if (enabled) {
+        await KillSwitchResource.enableKillSwitch(type);
+      } else {
+        await KillSwitchResource.disableKillSwitch(type);
+      }
+      return res.status(200).json({ success: true });
+
+    default:
+      return apiError(req, res, {
+        status_code: 405,
+        api_error: {
+          type: "method_not_supported_error",
+          message: "The method passed is not supported, GET is expected.",
+        },
+      });
+  }
+}
+
+export default withSessionAuthentication(handler);
diff --git a/front/pages/poke/kill.tsx b/front/pages/poke/kill.tsx
new file mode 100644
index 000000000000..e236dd21609d
--- /dev/null
+++ b/front/pages/poke/kill.tsx
@@ -0,0 +1,144 @@
+import { SliderToggle, Spinner } from "@dust-tt/sparkle";
+import { useSendNotification } from "@dust-tt/sparkle";
+import React, { useState } from "react";
+import { useSWRConfig } from "swr/_internal";
+
+import PokeNavbar from "@app/components/poke/PokeNavbar";
+import { withSuperUserAuthRequirements } from "@app/lib/iam/session";
+import type { KillSwitchType } from "@app/lib/poke/types";
+import { usePokeKillSwitches } from "@app/poke/swr/kill";
+
+export const getServerSideProps = withSuperUserAuthRequirements<object>(
+  async () => {
+    return {
+      props: {},
+    };
+  }
+);
+
+const KillPage = () => {
+  const { killSwitches, isKillSwitchesLoading } = usePokeKillSwitches();
+  const [loading, setLoading] = useState(false);
+  const { mutate } = useSWRConfig();
+  const sendNotification = useSendNotification();
+
+  const killSwitchMap: Record<
+    KillSwitchType,
+    {
+      title: string;
+      description: string;
+    }
+  > = {
+    save_agent_configurations: {
+      title: "Agent Configurations",
+      description: "Disable saving of agent configurations",
+    },
+    save_data_source_views: {
+      title: "Data Source Views",
+      description: "Disable saving of data source views",
+    },
+  };
+
+  async function toggleKillSwitch(killSwitch: KillSwitchType) {
+    if (loading) {
+      return;
+    }
+
+    setLoading(true);
+
+    const isEnabled = killSwitches.includes(killSwitch);
+
+    if (!isEnabled) {
+      if (
+        !window.confirm(
+          `Are you sure you want to enable the ${killSwitchMap[killSwitch].title} kill switch?`
+        )
+      ) {
+        setLoading(false);
+        return;
+      }
+    }
+
+    const res = await fetch(`/api/poke/kill`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        enabled: !isEnabled,
+        type: killSwitch,
+      }),
+    });
+
+    if (res.ok) {
+      await mutate("/api/poke/kill");
+      sendNotification({
+        title: "Kill switch updated",
+        description: `Kill switch ${killSwitch} updated`,
+        type: "success",
+      });
+    } else {
+      const errorData = await res.json();
+      console.error(errorData);
+      sendNotification({
+        title: "Error updating kill switch",
+        description: `Error: ${errorData.error.message}`,
+        type: "error",
+      });
+    }
+
+    setLoading(false);
+  }
+
+  return (
+    <div className="min-h-screen bg-structure-50">
+      <PokeNavbar />
+
+      <main className="mx-auto max-w-7xl px-4 sm:px-6 lg:px-8">
+        <div className="py-12">
+          <div className="text-center">
+            <h1 className="text-3xl font-bold tracking-tight text-gray-900">
+              System Kill Switches
+            </h1>
+            <p className="mt-2 text-sm text-gray-600">
+              Control critical system functionalities
+            </p>
+          </div>
+
+          {isKillSwitchesLoading ? (
+            <Spinner />
+          ) : (
+            <div className="mx-auto mt-12 max-w-xl">
+              <div className="space-y-6 rounded-lg bg-white p-6 shadow-sm">
+                {Object.entries(killSwitchMap).map(([key, value]) => {
+                  return (
+                    <div
+                      className="flex items-center justify-between border-b border-gray-100 py-4"
+                      key={key}
+                    >
+                      <div>
+                        <h3 className="text-sm font-medium text-gray-900">
+                          {value.title}
+                        </h3>
+                        <p className="text-sm text-gray-500">
+                          {value.description}
+                        </p>
+                      </div>
+                      <SliderToggle
+                        onClick={() => toggleKillSwitch(key as KillSwitchType)}
+                        selected={killSwitches.includes(key as KillSwitchType)}
+                        disabled={loading}
+                      />
+                    </div>
+                  );
+                })}
+              </div>
+            </div>
+          )}
+        </div>
+      </main>
+    </div>
+  );
+};
+
+export default KillPage;
diff --git a/front/poke/swr/kill.ts b/front/poke/swr/kill.ts
new file mode 100644
index 000000000000..202f83c7aa38
--- /dev/null
+++ b/front/poke/swr/kill.ts
@@ -0,0 +1,19 @@
+import { useMemo } from "react";
+import type { Fetcher } from "swr";
+import useSWR from "swr";
+
+import { fetcher } from "@app/lib/swr/swr";
+import type { GetKillSwitchesResponseBody } from "@app/pages/api/poke/kill";
+
+export function usePokeKillSwitches() {
+  const killSwitchesFetcher: Fetcher<GetKillSwitchesResponseBody> = fetcher;
+
+  const { data, error, mutate } = useSWR("/api/poke/kill", killSwitchesFetcher);
+
+  return {
+    killSwitches: useMemo(() => (data ? data.killSwitches : []), [data]),
+    isKillSwitchesLoading: !error && !data,
+    isKillSwitchesError: error,
+    mutateKillSwitches: mutate,
+  };
+}