From cdafe3c454ea239d717020e63e418af319d90028 Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <ryan@acorn.io>
Date: Tue, 14 Jan 2025 16:19:33 -0600
Subject: [PATCH] enhance: update UX for tool authentication

---
 .../agent/shared/ToolAuthenticationStatus.tsx | 106 +++++++++++-------
 .../composed/ConfirmationDialog.tsx           |  11 +-
 2 files changed, 77 insertions(+), 40 deletions(-)

diff --git a/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx b/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx
index 8c603bee8..90acb8df4 100644
--- a/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx
+++ b/ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx
@@ -1,4 +1,4 @@
-import { ScanEyeIcon, UserRoundCheckIcon } from "lucide-react";
+import { LockIcon, LockOpenIcon } from "lucide-react";
 
 import { ToolInfo } from "~/lib/model/agents";
 import { AssistantNamespace } from "~/lib/model/assistants";
@@ -9,13 +9,7 @@ import { useToolReference } from "~/components/agent/ToolEntry";
 import { ToolAuthenticationDialog } from "~/components/agent/shared/ToolAuthenticationDialog";
 import { ConfirmationDialog } from "~/components/composed/ConfirmationDialog";
 import { Button } from "~/components/ui/button";
-import {
-	DropdownMenu,
-	DropdownMenuContent,
-	DropdownMenuItem,
-	DropdownMenuLabel,
-	DropdownMenuTrigger,
-} from "~/components/ui/dropdown-menu";
+import { DialogDescription } from "~/components/ui/dialog";
 import {
 	Tooltip,
 	TooltipContent,
@@ -49,11 +43,11 @@ export function ToolAuthenticationStatus({
 
 	const { credentialNames, authorized } = toolInfo?.[tool] ?? {};
 
-	const { interceptAsync, dialogProps } = useConfirmationDialog();
+	const deleteConfirm = useConfirmationDialog();
+	const authorizeConfirm = useConfirmationDialog();
 
-	const handleAuthorize = async () => {
-		authorize.execute(namespace, entityId, [tool]);
-	};
+	const handleAuthorize = () =>
+		authorize.executeAsync(namespace, entityId, [tool]);
 
 	const handleDeauthorize = async () => {
 		if (!toolInfo) return;
@@ -93,43 +87,77 @@ export function ToolAuthenticationStatus({
 			</Tooltip>
 		);
 
+	const handleClick = () => {
+		if (authorized) {
+			deleteConfirm.interceptAsync(handleDeauthorize);
+		} else {
+			authorizeConfirm.interceptAsync(handleAuthorize);
+		}
+	};
+
 	if (!credentialNames?.length) return null;
 
 	return (
 		<>
 			<Tooltip>
-				<DropdownMenu>
-					<DropdownMenuTrigger asChild>
-						<TooltipTrigger asChild>
-							<Button size="icon" variant="ghost" loading={loading}>
-								{authorized ? <UserRoundCheckIcon /> : <ScanEyeIcon />}
-							</Button>
-						</TooltipTrigger>
-					</DropdownMenuTrigger>
-
-					<DropdownMenuContent side="right" align="start">
-						<DropdownMenuLabel>
-							{authorized ? "Authorized" : "Unauthorized"}
-						</DropdownMenuLabel>
+				<TooltipContent className="max-w-xs">
+					{authorized ? (
+						<>
+							<b>Authorized: </b>
+							Tool will use pre-authenticated credentials for each thread.
+						</>
+					) : (
+						<>
+							<b>Unauthorized: </b>
+							Tool will require user authentication for each thread.
+						</>
+					)}
+				</TooltipContent>
 
+				<TooltipTrigger asChild>
+					<Button
+						size="icon"
+						variant="ghost"
+						loading={loading}
+						onClick={handleClick}
+					>
 						{authorized ? (
-							<DropdownMenuItem
-								variant="destructive"
-								onClick={() => interceptAsync(handleDeauthorize)}
-							>
-								Remove Authorization
-							</DropdownMenuItem>
+							<LockOpenIcon className="text-success" />
 						) : (
-							<DropdownMenuItem onClick={handleAuthorize}>
-								Authorize Tool
-							</DropdownMenuItem>
+							<LockIcon />
 						)}
-					</DropdownMenuContent>
-				</DropdownMenu>
-
-				<TooltipContent>Authorization Status</TooltipContent>
+					</Button>
+				</TooltipTrigger>
 			</Tooltip>
 
+			<ConfirmationDialog
+				{...authorizeConfirm.dialogProps}
+				title={
+					<span className="flex items-center gap-2">
+						{icon}
+						Pre-Authenticate {label}?
+					</span>
+				}
+				content={
+					<>
+						<DialogDescription>
+							{label} is currently not authenticated. Users will be prompted for
+							authentication when using this tool in a new thread.
+						</DialogDescription>
+
+						<DialogDescription>
+							You can pre-authenticate {label} to allow users to use this tool
+							without authentication.
+						</DialogDescription>
+					</>
+				}
+				confirmProps={{
+					children: `Authenticate ${label}`,
+					loading: authorize.isLoading,
+					disabled: authorize.isLoading,
+				}}
+			/>
+
 			<ToolAuthenticationDialog
 				tool={tool}
 				entityId={entityId}
@@ -138,7 +166,7 @@ export function ToolAuthenticationStatus({
 			/>
 
 			<ConfirmationDialog
-				{...dialogProps}
+				{...deleteConfirm.dialogProps}
 				title={
 					<span className="flex items-center gap-2">
 						<span>{icon}</span>
diff --git a/ui/admin/app/components/composed/ConfirmationDialog.tsx b/ui/admin/app/components/composed/ConfirmationDialog.tsx
index 56b2eb5d4..e22d675f6 100644
--- a/ui/admin/app/components/composed/ConfirmationDialog.tsx
+++ b/ui/admin/app/components/composed/ConfirmationDialog.tsx
@@ -7,6 +7,7 @@ import {
 	DialogContent,
 	DialogDescription,
 	DialogFooter,
+	DialogHeader,
 	DialogTitle,
 	DialogTrigger,
 } from "~/components/ui/dialog";
@@ -15,6 +16,7 @@ export type ConfirmationDialogProps = ComponentProps<typeof Dialog> & {
 	children?: ReactNode;
 	title: ReactNode;
 	description?: ReactNode;
+	content?: ReactNode;
 	onConfirm: (e: React.MouseEvent<HTMLButtonElement>) => void;
 	onCancel?: (e: React.MouseEvent<HTMLButtonElement>) => void;
 	confirmProps?: Omit<Partial<ComponentProps<typeof Button>>, "onClick">;
@@ -25,6 +27,7 @@ export function ConfirmationDialog({
 	children,
 	title,
 	description,
+	content,
 	onConfirm,
 	onCancel,
 	confirmProps,
@@ -36,8 +39,14 @@ export function ConfirmationDialog({
 			{children && <DialogTrigger asChild>{children}</DialogTrigger>}
 
 			<DialogContent onClick={(e) => e.stopPropagation()}>
-				<DialogTitle>{title}</DialogTitle>
+				<DialogHeader>
+					<DialogTitle>{title}</DialogTitle>
+				</DialogHeader>
+
 				<DialogDescription>{description}</DialogDescription>
+
+				{content}
+
 				<DialogFooter>
 					<DialogClose onClick={onCancel} asChild>
 						<Button variant="secondary">Cancel</Button>