Skip to content

Commit

Permalink
feat: add kill switches in poke (#9336)
Browse files Browse the repository at this point in the history
Co-authored-by: Henry Fontanier <[email protected]>
  • Loading branch information
fontanierh and Henry Fontanier authored Dec 13, 2024
1 parent b6fbf41 commit fc9ece6
Show file tree
Hide file tree
Showing 8 changed files with 379 additions and 0 deletions.
1 change: 1 addition & 0 deletions front/components/poke/PokeNavbar.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 />
Expand Down
8 changes: 8 additions & 0 deletions front/lib/poke/types.ts
Original file line number Diff line number Diff line change
@@ -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);
}
64 changes: 64 additions & 0 deletions front/lib/resources/kill_switch_resource.ts
Original file line number Diff line number Diff line change
@@ -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);
}
}
43 changes: 43 additions & 0 deletions front/lib/resources/storage/models/kill_switches.ts
Original file line number Diff line number Diff line change
@@ -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"] }],
}
);
9 changes: 9 additions & 0 deletions front/migrations/db/migration_129.sql
Original file line number Diff line number Diff line change
@@ -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");
91 changes: 91 additions & 0 deletions front/pages/api/poke/kill.ts
Original file line number Diff line number Diff line change
@@ -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);
144 changes: 144 additions & 0 deletions front/pages/poke/kill.tsx
Original file line number Diff line number Diff line change
@@ -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;
19 changes: 19 additions & 0 deletions front/poke/swr/kill.ts
Original file line number Diff line number Diff line change
@@ -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,
};
}

0 comments on commit fc9ece6

Please sign in to comment.