-
Notifications
You must be signed in to change notification settings - Fork 111
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add kill switches in poke (#9336)
Co-authored-by: Henry Fontanier <[email protected]>
- Loading branch information
1 parent
b6fbf41
commit fc9ece6
Showing
8 changed files
with
379 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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"] }], | ||
} | ||
); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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"); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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, | ||
}; | ||
} |