enhance: improve tool authentication dialog UX

ui/admin/app/components/agent/shared/ToolAuthenticationDialog.tsx

@@ -1,9 +1,11 @@
-import { useCallback, useState } from "react";
+import { CheckIcon, CircleAlert } from "lucide-react";
+import { useEffect, useMemo, useState } from "react";
 
-import { ChatEvent } from "~/lib/model/chatEvents";
+import { ThreadsService } from "~/lib/service/api/threadsService";
 
 import { useToolReference } from "~/components/agent/ToolEntry";
-import { Chat, ChatProvider } from "~/components/chat";
+import { PromptAuthForm } from "~/components/chat/Message";
+import { LoadingSpinner } from "~/components/ui/LoadingSpinner";
 import { Button } from "~/components/ui/button";
 import {
 	Dialog,
@@ -14,35 +16,59 @@ import {
 	DialogHeader,
 	DialogTitle,
 } from "~/components/ui/dialog";
+import { Link } from "~/components/ui/link";
+import { useMessageStream } from "~/hooks/messages/useMessageSource";
 
 type AgentAuthenticationDialogProps = {
 	threadId: Nullish<string>;
 	onComplete: () => void;
-	entityId: string;
 	tool: string;
 };
 
 export function ToolAuthenticationDialog({
 	onComplete,
 	threadId,
-	entityId,
 	tool,
 }: AgentAuthenticationDialogProps) {
-	const [done, setDone] = useState(false);
-	const handleDone = useCallback(() => setDone(true), []);
-
 	const { icon, label } = useToolReference(tool);
 
-	const onRunEvent = useCallback(
-		({ content }: ChatEvent) => {
-			if (content === "DONE") handleDone();
-		},
-		[handleDone]
+	const source = useMemo(
+		() => (threadId ? ThreadsService.getThreadEventSource(threadId) : null),
+		[threadId]
 	);
+	const { messages: _messages } = useMessageStream(source);
+
+	type ItemState = {
+		isLoading?: boolean;
+		isError?: boolean;
+		isDone?: boolean;
+	};
+
+	const [map, setMap] = useState<Record<number, ItemState>>({});
+	const updateItem = (id: number, state: Partial<ItemState>) =>
+		setMap((prev) => ({ ...prev, [id]: { ...prev[id], ...state } }));
+
+	const messages = useMemo(
+		() => _messages.filter((m) => m.prompt || m.error || m.text === "DONE"),
+		[_messages]
+	);
+
+	useEffect(() => {
+		// any time a message is added, prevent the last message from being loading
+		const isError = messages.at(-1)?.error;
+
+		const i = messages.length - 2;
+		setMap((prev) => ({
+			...prev,
+			[i]: { isLoading: false, isDone: !isError, isError },
+		}));
+	}, [messages]);
+
+	const done = messages.at(-1)?.text === "DONE";
 
 	return (
 		<Dialog open={!!threadId} onOpenChange={onComplete}>
-			<DialogContent>
+			<DialogContent className="max-w-sm">
 				<DialogHeader>
 					<DialogTitle className="flex items-center gap-2">
 						{icon} <span>Authorize {label}</span>
@@ -51,25 +77,94 @@ export function ToolAuthenticationDialog({
 					<DialogDescription hidden={done}></DialogDescription>
 				</DialogHeader>
 
-				{done ? (
-					<DialogDescription>
-						{label} has successfully been authorized. You may now close this
-						modal.
-					</DialogDescription>
-				) : (
-					<ChatProvider
-						threadId={threadId}
-						id={entityId}
-						readOnly
-						onRunEvent={onRunEvent}
-					>
-						<Chat
-							classNames={{
-								messagePane: { messageList: "px-0" },
-							}}
-						/>
-					</ChatProvider>
-				)}
+				<div className="flex w-full items-center justify-center [&_svg]:size-4">
+					{!messages.length ? (
+						<div className="flex items-center gap-2">
+							<LoadingSpinner /> Loading...
+						</div>
+					) : (
+						<div className="flex flex-col gap-2">
+							{messages.map((message, index) => {
+								if (message.error) {
+									return (
+										<p
+											className="flex items-center gap-2 text-destructive"
+											key={index}
+										>
+											<CircleAlert /> Error: {message.text}
+										</p>
+									);
+								}
+
+								if (message.text === "DONE") {
+									return (
+										<p key={index} className="flex items-center gap-2">
+											<CheckIcon className="text-success" />
+											Done
+										</p>
+									);
+								}
+
+								if (message.prompt) {
+									if (map[index]?.isDone) {
+										return (
+											<p key={index} className="flex items-center gap-2">
+												<CheckIcon className="text-success" />
+												Authentication Successful
+											</p>
+										);
+									}
+
+									if (map[index]?.isLoading) {
+										return (
+											<p key={index} className="flex items-center gap-2">
+												<LoadingSpinner /> Authentication Processing
+											</p>
+										);
+									}
+
+									if (message.prompt.metadata?.authURL) {
+										return (
+											<p key={index} className="flex items-center gap-2">
+												<CircleAlert />
+												<span>
+													Authentication Required{" "}
+													<Link
+														target="_blank"
+														rel="noreferrer"
+														to={message.prompt.metadata.authURL}
+														onClick={() =>
+															updateItem(index, { isLoading: true })
+														}
+													>
+														Click Here
+													</Link>
+												</span>
+											</p>
+										);
+									}
+
+									if (message.prompt.fields) {
+										return (
+											<div key={index} className="flex flex-col gap-2">
+												<p className="flex items-center gap-2">
+													<CircleAlert />
+													Authentication Required
+												</p>
+												<PromptAuthForm
+													prompt={message.prompt}
+													onSubmit={() =>
+														updateItem(index, { isLoading: true })
+													}
+												/>
+											</div>
+										);
+									}
+								}
+							})}
+						</div>
+					)}
+				</div>
 
 				<DialogFooter>
 					<DialogClose asChild>

ui/admin/app/components/agent/shared/ToolAuthenticationStatus.tsx

@@ -160,7 +160,6 @@ export function ToolAuthenticationStatus({
 
 			<ToolAuthenticationDialog
 				tool={tool}
-				entityId={entityId}
 				threadId={threadId}
 				onComplete={handleAuthorizeComplete}
 			/>