From dde0a225ab2d3a6b239025d6d4dd6750debd2048 Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <46546486+ryanhopperlowe@users.noreply.github.com>
Date: Mon, 2 Dec 2024 13:51:34 -0600
Subject: [PATCH] feat: allow users to trigger workflow authentication in the
 chat interface (#725)

---
 ui/admin/app/components/chat/Message.tsx      | 80 +++++++++++++------
 ui/admin/app/components/workflow/Workflow.tsx | 24 +++++-
 ui/admin/app/lib/model/chatEvents.ts          | 23 ++----
 ui/admin/app/lib/routers/apiRoutes.ts         |  2 +
 .../app/lib/service/api/workflowService.ts    | 21 ++++-
 .../app/routes/_auth.workflows.$workflow.tsx  | 28 ++++---
 6 files changed, 125 insertions(+), 53 deletions(-)

diff --git a/ui/admin/app/components/chat/Message.tsx b/ui/admin/app/components/chat/Message.tsx
index 8bd6b27f..1cdcde7f 100644
--- a/ui/admin/app/components/chat/Message.tsx
+++ b/ui/admin/app/components/chat/Message.tsx
@@ -21,6 +21,7 @@ import { Button } from "~/components/ui/button";
 import {
     Dialog,
     DialogContent,
+    DialogDescription,
     DialogHeader,
     DialogTitle,
     DialogTrigger,
@@ -84,7 +85,7 @@ export const Message = React.memo(({ message }: MessageProps) => {
                             />
                         )}
 
-                        {message.prompt?.metadata ? (
+                        {message.prompt ? (
                             <PromptMessage prompt={message.prompt} />
                         ) : (
                             <Markdown
@@ -141,26 +142,51 @@ Message.displayName = "Message";
 
 function PromptMessage({ prompt }: { prompt: AuthPrompt }) {
     const [open, setOpen] = useState(false);
+    const [isSubmitted, setIsSubmitted] = useState(false);
 
-    if (!prompt.metadata) return null;
+    const getMessage = () => {
+        if (prompt.metadata?.authURL || prompt.metadata?.authType)
+            return `${prompt.metadata.category || "Tool call"} requires Authentication`;
+
+        return prompt.message;
+    };
+
+    const getCtaText = () => {
+        if (prompt.metadata?.authURL || prompt.metadata?.authType)
+            return ["Authenticate", prompt.metadata.category]
+                .filter(Boolean)
+                .join(" with ");
+
+        return "Submit Parameters";
+    };
+
+    const getSubmittedText = () => {
+        if (prompt.metadata?.authURL || prompt.metadata?.authType)
+            return "Authenticated";
+
+        return "Parameters Submitted";
+    };
+
+    if (isSubmitted) {
+        return (
+            <div className="flex-auto flex flex-col flex-wrap gap-2 w-fit">
+                <TypographyP className="min-w-fit">
+                    {getSubmittedText()}
+                </TypographyP>
+            </div>
+        );
+    }
 
     return (
         <div className="flex-auto flex flex-col flex-wrap gap-2 w-fit">
-            <TypographyP className="min-w-fit">
-                <b>
-                    {[prompt.metadata?.category, prompt.name]
-                        .filter(Boolean)
-                        .join(" - ")}
-                </b>
-                {": "}
-                Tool Call requires authentication
-            </TypographyP>
-
-            {prompt.metadata.authType === "oauth" && (
+            <TypographyP className="min-w-fit">{getMessage()}</TypographyP>
+
+            {prompt.metadata?.authURL && (
                 <Link
                     as="button"
                     rel="noreferrer"
                     target="_blank"
+                    onClick={() => setIsSubmitted(true)}
                     to={prompt.metadata.authURL}
                 >
                     <ToolIcon
@@ -169,37 +195,41 @@ function PromptMessage({ prompt }: { prompt: AuthPrompt }) {
                         name={prompt.name}
                         disableTooltip
                     />
-                    Authenticate with {prompt.metadata.category}
+
+                    {getCtaText()}
                 </Link>
             )}
 
-            {prompt.metadata.authType === "basic" && prompt.fields && (
+            {prompt.fields && (
                 <Dialog open={open} onOpenChange={setOpen}>
-                    <DialogTrigger asChild>
+                    <DialogTrigger disabled={isSubmitted} asChild>
                         <Button
                             startContent={
                                 <ToolIcon
-                                    icon={prompt.metadata.icon}
-                                    category={prompt.metadata.category}
+                                    icon={prompt.metadata?.icon}
+                                    category={prompt.metadata?.category}
                                     name={prompt.name}
                                     disableTooltip
                                 />
                             }
                         >
-                            Authenticate with {prompt.metadata.category}
+                            {getCtaText()}
                         </Button>
                     </DialogTrigger>
 
                     <DialogContent>
                         <DialogHeader>
-                            <DialogTitle>
-                                Authenticate with {prompt.metadata.category}
-                            </DialogTitle>
+                            <DialogTitle>{getCtaText()}</DialogTitle>
                         </DialogHeader>
 
+                        <DialogDescription>{prompt.message}</DialogDescription>
+
                         <PromptAuthForm
                             prompt={prompt}
-                            onSuccess={() => setOpen(false)}
+                            onSuccess={() => {
+                                setOpen(false);
+                                setIsSubmitted(true);
+                            }}
                         />
                     </DialogContent>
                 </Dialog>
@@ -242,7 +272,7 @@ function PromptAuthForm({
                         control={form.control}
                         name={field}
                         label={field}
-                        type={field.includes("password") ? "password" : "text"}
+                        type={prompt.sensitive ? "password" : "text"}
                     />
                 ))}
 
@@ -251,7 +281,7 @@ function PromptAuthForm({
                     loading={authenticate.isLoading}
                     type="submit"
                 >
-                    Authenticate
+                    Submit
                 </Button>
             </form>
         </Form>
diff --git a/ui/admin/app/components/workflow/Workflow.tsx b/ui/admin/app/components/workflow/Workflow.tsx
index da65a49a..baf56090 100644
--- a/ui/admin/app/components/workflow/Workflow.tsx
+++ b/ui/admin/app/components/workflow/Workflow.tsx
@@ -2,6 +2,7 @@ import {
     KeyIcon,
     Library,
     List,
+    LockIcon,
     PuzzleIcon,
     Variable,
     WrenchIcon,
@@ -9,12 +10,14 @@ import {
 import { useCallback, useState } from "react";
 
 import { Workflow as WorkflowType } from "~/lib/model/workflows";
+import { WorkflowService } from "~/lib/service/api/workflowService";
 import { cn } from "~/lib/utils";
 
 import { TypographyH4, TypographyP } from "~/components/Typography";
 import { AgentForm } from "~/components/agent";
 import { AgentKnowledgePanel } from "~/components/knowledge";
 import { BasicToolForm } from "~/components/tools/BasicToolForm";
+import { Button } from "~/components/ui/button";
 import { CardDescription } from "~/components/ui/card";
 import { ScrollArea } from "~/components/ui/scroll-area";
 import { ParamsForm } from "~/components/workflow/ParamsForm";
@@ -25,10 +28,12 @@ import {
 } from "~/components/workflow/WorkflowContext";
 import { WorkflowEnvForm } from "~/components/workflow/WorkflowEnvForm";
 import { StepsForm } from "~/components/workflow/steps/StepsForm";
+import { useAsync } from "~/hooks/useAsync";
 import { useDebounce } from "~/hooks/useDebounce";
 
 type WorkflowProps = {
     workflow: WorkflowType;
+    onPersistThreadId: (threadId: string) => void;
     className?: string;
 };
 
@@ -40,7 +45,7 @@ export function Workflow(props: WorkflowProps) {
     );
 }
 
-function WorkflowContent({ className }: WorkflowProps) {
+function WorkflowContent({ className, onPersistThreadId }: WorkflowProps) {
     const { workflow, updateWorkflow, isUpdating, lastUpdated } = useWorkflow();
 
     const [workflowUpdates, setWorkflowUpdates] = useState(workflow);
@@ -62,6 +67,10 @@ function WorkflowContent({ className }: WorkflowProps) {
 
     const debouncedSetWorkflowInfo = useDebounce(partialSetWorkflow, 1000);
 
+    const authenticate = useAsync(WorkflowService.authenticateWorkflow, {
+        onSuccess: ({ threadId }) => onPersistThreadId(threadId),
+    });
+
     return (
         <div className="h-full flex flex-col">
             <ScrollArea className={cn("h-full", className)}>
@@ -168,7 +177,7 @@ function WorkflowContent({ className }: WorkflowProps) {
                 </div>
             </ScrollArea>
 
-            <footer className="flex justify-between items-center p-4 gap-4 text-muted-foreground">
+            <footer className="flex justify-between items-center p-4 gap-4 border-t text-muted-foreground">
                 {isUpdating ? (
                     <TypographyP>Saving...</TypographyP>
                 ) : lastUpdated ? (
@@ -176,6 +185,17 @@ function WorkflowContent({ className }: WorkflowProps) {
                 ) : (
                     <div />
                 )}
+
+                <div className="flex items-center gap-2">
+                    <Button
+                        onClick={() => authenticate.execute(workflow.id)}
+                        loading={authenticate.isLoading}
+                        disabled={authenticate.isLoading}
+                        startContent={<LockIcon />}
+                    >
+                        Authenticate
+                    </Button>
+                </div>
             </footer>
         </div>
     );
diff --git a/ui/admin/app/lib/model/chatEvents.ts b/ui/admin/app/lib/model/chatEvents.ts
index 2ac485d7..1722137e 100644
--- a/ui/admin/app/lib/model/chatEvents.ts
+++ b/ui/admin/app/lib/model/chatEvents.ts
@@ -13,20 +13,13 @@ export type ToolCall = {
     };
 };
 
-type PromptAuthMetaBase = {
-    category: string;
-    icon: string;
-    toolContext: string;
-    toolDisplayName: string;
-};
-
-type PromptOAuthMeta = PromptAuthMetaBase & {
-    authType: "oauth";
-    authURL: string;
-};
-
-type PromptAuthBasicMeta = PromptAuthMetaBase & {
-    authType: "basic";
+type PromptAuthMeta = {
+    authURL?: string;
+    category?: string;
+    icon?: string;
+    toolContext?: string;
+    toolDisplayName?: string;
+    authType: "oauth" | "basic";
 };
 
 export type AuthPrompt = {
@@ -36,7 +29,7 @@ export type AuthPrompt = {
     message: string;
     fields?: string[];
     sensitive?: boolean;
-    metadata?: PromptOAuthMeta | PromptAuthBasicMeta;
+    metadata?: PromptAuthMeta;
 };
 
 // note(ryanhopperlowe) renaming this to ChatEvent to differentiate itself specifically for a chat with an agent
diff --git a/ui/admin/app/lib/routers/apiRoutes.ts b/ui/admin/app/lib/routers/apiRoutes.ts
index 7f3e08b9..ae2d1b6a 100644
--- a/ui/admin/app/lib/routers/apiRoutes.ts
+++ b/ui/admin/app/lib/routers/apiRoutes.ts
@@ -109,6 +109,8 @@ export const ApiRoutes = {
             buildUrl(`/workflows/${workflowId}/files/${fileName}`),
         deleteKnowledge: (workflowId: string, fileName: string) =>
             buildUrl(`/workflows/${workflowId}/files/${fileName}`),
+        authenticate: (workflowId: string) =>
+            buildUrl(`/workflows/${workflowId}/authenticate`),
     },
     threads: {
         base: () => buildUrl("/threads"),
diff --git a/ui/admin/app/lib/service/api/workflowService.ts b/ui/admin/app/lib/service/api/workflowService.ts
index 3420ca65..deaa5f24 100644
--- a/ui/admin/app/lib/service/api/workflowService.ts
+++ b/ui/admin/app/lib/service/api/workflowService.ts
@@ -4,7 +4,7 @@ import {
     Workflow,
 } from "~/lib/model/workflows";
 import { ApiRoutes, revalidateWhere } from "~/lib/routers/apiRoutes";
-import { request } from "~/lib/service/api/primitives";
+import { ResponseHeaders, request } from "~/lib/service/api/primitives";
 
 async function getWorkflows() {
     const res = await request<{ items: Workflow[] }>({
@@ -69,6 +69,24 @@ async function deleteWorkflow(id: string) {
 const revalidateWorkflows = () =>
     revalidateWhere((url) => url.includes(ApiRoutes.workflows.base().path));
 
+async function authenticateWorkflow(workflowId: string) {
+    const response = await request<ReadableStream>({
+        url: ApiRoutes.workflows.authenticate(workflowId).url,
+        method: "POST",
+        headers: { Accept: "text/event-stream" },
+        responseType: "stream",
+        errorMessage: "Failed to invoke agenticate workflow",
+    });
+
+    const reader = response.data
+        ?.pipeThrough(new TextDecoderStream())
+        .getReader();
+
+    const threadId = response.headers[ResponseHeaders.ThreadId] as string;
+
+    return { reader, threadId };
+}
+
 export const WorkflowService = {
     getWorkflows,
     getWorkflowById,
@@ -76,4 +94,5 @@ export const WorkflowService = {
     updateWorkflow,
     deleteWorkflow,
     revalidateWorkflows,
+    authenticateWorkflow,
 };
diff --git a/ui/admin/app/routes/_auth.workflows.$workflow.tsx b/ui/admin/app/routes/_auth.workflows.$workflow.tsx
index 3785c1b1..f66784ed 100644
--- a/ui/admin/app/routes/_auth.workflows.$workflow.tsx
+++ b/ui/admin/app/routes/_auth.workflows.$workflow.tsx
@@ -4,6 +4,7 @@ import {
     useLoaderData,
     useNavigate,
 } from "@remix-run/react";
+import { useCallback } from "react";
 import { $path } from "remix-routes";
 import { preload } from "swr";
 
@@ -48,28 +49,35 @@ export default function ChatAgent() {
 
     const navigate = useNavigate();
 
+    const onPersistThreadId = useCallback(
+        (threadId: string) =>
+            navigate(
+                $path(
+                    "/workflows/:workflow",
+                    { workflow: workflow.id },
+                    { threadId }
+                )
+            ),
+        [navigate, workflow.id]
+    );
+
     return (
         <div className="h-full flex flex-col overflow-hidden relative">
             <ChatProvider
                 id={workflow.id}
                 mode="workflow"
                 threadId={threadId}
-                onCreateThreadId={(threadId) =>
-                    navigate(
-                        $path(
-                            "/workflows/:workflow",
-                            { workflow: workflow.id },
-                            { threadId }
-                        )
-                    )
-                }
+                onCreateThreadId={onPersistThreadId}
             >
                 <ResizablePanelGroup
                     direction="horizontal"
                     className="flex-auto"
                 >
                     <ResizablePanel className="">
-                        <Workflow workflow={workflow} />
+                        <Workflow
+                            workflow={workflow}
+                            onPersistThreadId={onPersistThreadId}
+                        />
                     </ResizablePanel>
                     <ResizableHandle withHandle />
                     <ResizablePanel>