Fix tool execution bug

frontend/package.json

@@ -60,6 +60,7 @@
     "sonner": "^1.7.3",
     "tailwind-merge": "^2.6.0",
     "tailwindcss-animate": "^1.0.7",
+    "usehooks-ts": "^3.1.1",
     "zod": "^3.24.1",
     "zustand": "^5.0.3"
   },

frontend/src/app/playground/chat-input.tsx

@@ -1,69 +1,126 @@
 import { Button } from "@/components/ui/button";
 import { Textarea } from "@/components/ui/textarea";
 import { toast } from "sonner";
-import { ArrowRight } from "lucide-react";
+import { ArrowUp, StopCircle } from "lucide-react";
+import { useCallback, memo } from "react";
+import { UseChatHelpers } from "ai/react";
 
 interface ChatInputProps {
   input: string;
   handleInputChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
   handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void;
   status: string;
   linkedAccountOwnerId: string | null;
+  stop?: () => void;
+  setMessages?: UseChatHelpers["setMessages"];
 }
 
+function PureStopButton({
+  stop,
+  setMessages,
+}: {
+  stop: () => void;
+  setMessages: UseChatHelpers["setMessages"];
+}) {
+  return (
+    <Button
+      type="button"
+      data-testid="stop-button"
+      className="rounded-full p-1.5 h-fit border dark:border-zinc-600"
+      onClick={(event) => {
+        event.preventDefault();
+        stop();
+        setMessages((messages) => messages);
+      }}
+    >
+      <StopCircle size={14} />
+    </Button>
+  );
+}
+
+const StopButton = memo(PureStopButton);
+
+function PureSendButton({ input }: { input: string }) {
+  return (
+    <Button
+      data-testid="send-button"
+      className="rounded-full p-1.5 h-fit border dark:border-zinc-600"
+      disabled={input.length === 0}
+    >
+      <ArrowUp size={14} />
+    </Button>
+  );
+}
+
+const SendButton = memo(PureSendButton, (prevProps, nextProps) => {
+  if ((prevProps.input.length === 0) !== (nextProps.input.length === 0))
+    return false;
+  return true;
+});
+
 export function ChatInput({
   input,
   handleInputChange,
   handleSubmit,
   status,
   linkedAccountOwnerId,
+  stop,
+  setMessages,
 }: ChatInputProps) {
-  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-    if (e.key === "Enter" && !e.shiftKey) {
-      e.preventDefault();
-      if (linkedAccountOwnerId) {
-        const formEvent = new Event(
-          "submit",
-        ) as unknown as React.FormEvent<HTMLFormElement>;
-        handleSubmit(formEvent);
-      } else {
-        toast.error("Please select a linked account");
-      }
-    }
-  };
+  const submitForm = useCallback(
+    (e?: React.FormEvent<HTMLFormElement>) => {
+      e?.preventDefault();
+      handleSubmit(e as React.FormEvent<HTMLFormElement>);
+    },
+    [handleSubmit],
+  );
 
   return (
     <div className="pt-4 border-none relative -mt-6 z-10">
       <form
-        onSubmit={(event) => {
-          if (!linkedAccountOwnerId) {
-            toast.error("Please select a linked account");
-            return;
-          }
-          handleSubmit(event);
-        }}
+        onSubmit={submitForm}
         className="flex flex-col w-full max-w-3xl mx-auto"
       >
         <div className="flex flex-col items-start bg-white rounded-2xl border shadow-sm">
-          <Textarea
-            value={input}
-            placeholder="Ask me anything..."
-            onChange={handleInputChange}
-            onKeyDown={handleKeyDown}
-            className="flex-1 p-4 bg-transparent outline-none resize-none min-h-[4rem] border-none shadow-none focus-visible:ring-0 focus-visible:ring-offset-0 focus-visible:outline-none"
-            disabled={status !== "ready"}
-            rows={2}
-          />
-          <div className="flex flex-row justify-end w-full p-2">
-            <Button
-              type="submit"
-              variant="outline"
-              disabled={status !== "ready" || !linkedAccountOwnerId}
-              className="px-4 py-2 text-gray-500 hover:text-gray-900 disabled:opacity-50 flex items-center gap-2"
-            >
-              <span>Run</span>
-              <ArrowRight className="w-4 h-4" />
-            </Button>
+          <div className="relative w-full">
+            <Textarea
+              value={input}
+              placeholder="Ask me anything..."
+              onChange={handleInputChange}
+              onKeyDown={(event) => {
+                if (!linkedAccountOwnerId) {
+                  toast.error("Please select a linked account owner");
+                  return;
+                }
+                if (
+                  event.key === "Enter" &&
+                  !event.shiftKey &&
+                  linkedAccountOwnerId
+                ) {
+                  event.preventDefault();
+                  if (status === "submitted" || status === "streaming") {
+                    toast.error(
+                      "Please wait for the previous request to complete",
+                    );
+                    return;
+                  } else {
+                    submitForm();
+                  }
+                }
+              }}
+              className="flex-1 p-4 pr-20 bg-transparent outline-none resize-none min-h-[4rem] border-none shadow-none focus-visible:ring-0 focus-visible:ring-offset-0 focus-visible:outline-none"
+              disabled={status !== "ready"}
+              rows={2}
+            />
+            <div className="pointer-events-none">
+              <div className="absolute top-1/2 right-4 -translate-y-1/2 flex flex-row justify-end pointer-events-auto">
+                {status === "submitted" || status === "streaming" ? (
+                  <StopButton stop={stop!} setMessages={setMessages!} />
+                ) : (
+                  <SendButton input={input} />
+                )}
+              </div>
+            </div>
           </div>
         </div>
       </form>

frontend/src/app/playground/function-calling.tsx

@@ -1,16 +1,42 @@
+import { useToolExecution } from "@/hooks/use-tool-execution";
 import type { ToolInvocation } from "ai";
 
 type FunctionCallingProps = {
   toolInvocation: ToolInvocation;
+  linkedAccountOwnerId: string;
+  apiKey: string;
+  addToolResult: ({
+    toolCallId,
+    result,
+  }: {
+    toolCallId: string;
+    result: object;
+  }) => void;
 };
 
-export function FunctionCalling({ toolInvocation }: FunctionCallingProps) {
+export function FunctionCalling({
+  toolInvocation,
+  linkedAccountOwnerId,
+  apiKey,
+  addToolResult,
+}: FunctionCallingProps) {
   const { toolName } = toolInvocation;
 
+  const { error } = useToolExecution({
+    toolInvocation,
+    addToolResult,
+    linkedAccountOwnerId,
+    apiKey,
+  });
+
+  if (error) {
+    return <div>Function Calling Error: {error}</div>;
+  }
+
   return (
-    <div>
-      {/* TODO: Add tool call */}
-      Function Calling: {toolName}
+    <div className="flex items-center space-x-2 border border-gray-200 rounded-md p-2">
+      <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-gray-900"></div>
+      <span>Function Calling: {toolName}</span>
     </div>
   );
 }

frontend/src/app/playground/message.tsx

@@ -15,9 +15,21 @@ import { Markdown } from "./markdown";
 const PurePreviewMessage = ({
   message,
   isLoading,
+  linkedAccountOwnerId,
+  apiKey,
+  addToolResult,
 }: {
   message: UIMessage;
   isLoading: boolean;
+  linkedAccountOwnerId: string;
+  apiKey: string;
+  addToolResult: ({
+    toolCallId,
+    result,
+  }: {
+    toolCallId: string;
+    result: object;
+  }) => void;
 }) => {
   return (
     <AnimatePresence>
@@ -103,6 +115,9 @@ const PurePreviewMessage = ({
                     <FunctionCalling
                       key={`${toolCallId}-${state}`}
                       toolInvocation={toolInvocation}
+                      linkedAccountOwnerId={linkedAccountOwnerId}
+                      apiKey={apiKey}
+                      addToolResult={addToolResult}
                     />
                   );
                 }

frontend/src/app/playground/messages.tsx

@@ -9,9 +9,24 @@ import { Overview } from "./overview";
 interface MessagesProps {
   status: UseChatHelpers["status"];
   messages: Array<UIMessage>;
+  linkedAccountOwnerId: string;
+  apiKey: string;
+  addToolResult: ({
+    toolCallId,
+    result,
+  }: {
+    toolCallId: string;
+    result: object;
+  }) => void;
 }
 
-function PureMessages({ status, messages }: MessagesProps) {
+function PureMessages({
+  status,
+  messages,
+  linkedAccountOwnerId,
+  apiKey,
+  addToolResult,
+}: MessagesProps) {
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>(messages);
 
@@ -31,6 +46,9 @@ function PureMessages({ status, messages }: MessagesProps) {
             index === messages.length - 1 &&
             message.role === "assistant"
           }
+          linkedAccountOwnerId={linkedAccountOwnerId}
+          apiKey={apiKey}
+          addToolResult={addToolResult}
         />
       ))}
 

frontend/src/app/playground/page.tsx

@@ -3,7 +3,6 @@
 import { useMetaInfo } from "@/components/context/metainfo";
 import { useChat } from "@ai-sdk/react";
 import { useAgentStore } from "@/lib/store/agent";
-import { executeFunction, searchFunctions } from "@/lib/api/appfunction";
 import { SettingsSidebar } from "./playground-settings";
 import { ChatInput } from "./chat-input";
 import { Messages } from "./messages";
@@ -38,6 +37,7 @@ const Page = () => {
     status,
     addToolResult,
     setMessages,
+    stop,
   } = useChat({
     api: `${process.env.NEXT_PUBLIC_API_URL}/v1/agent/chat`,
     headers: {
@@ -54,51 +54,18 @@ const Page = () => {
     onFinish: (message) => {
       console.log(message);
     },
-    onToolCall: async ({ toolCall }) => {
-      // TODO: Human in the loop
-      console.log("tool call", toolCall);
-
-      let result;
-      // when the tool call name is "ACI_SEARCH_FUNCTIONS"
-      if (toolCall.toolName === "ACI_SEARCH_FUNCTIONS") {
-        result = await searchFunctions(
-          toolCall.args as Record<string, unknown>,
-          apiKey,
-        );
-        addToolResult({
-          toolCallId: toolCall.toolCallId,
-          result: result,
-        });
-      } else if (toolCall.toolName === "ACI_EXECUTE_FUNCTION") {
-        result = await executeFunction(
-          toolCall.toolName,
-          {
-            function_input: toolCall.args as Record<string, unknown>,
-            linked_account_owner_id: selectedLinkedAccountOwnerId,
-          },
-          apiKey,
-        );
-        addToolResult({
-          toolCallId: toolCall.toolCallId,
-          result: result,
-        });
-      } else {
-        result = await executeFunction(
-          toolCall.toolName,
-          {
-            function_input: toolCall.args as Record<string, unknown>,
-            linked_account_owner_id: selectedLinkedAccountOwnerId,
-          },
-          apiKey,
-        );
-        addToolResult({
-          toolCallId: toolCall.toolCallId,
-          result: result,
-        });
-      }
-    },
   });
 
+  const handleAddToolResult = ({
+    toolCallId,
+    result,
+  }: {
+    toolCallId: string;
+    result: object;
+  }) => {
+    addToolResult({ toolCallId, result });
+  };
+
   if (!activeProject) {
     console.warn("No active project");
     return <div>No project selected</div>;
@@ -109,9 +76,17 @@ const Page = () => {
       {/* Left part - Chat area */}
       <div className="flex-1 flex flex-col overflow-hidden">
         <BetaAlert />
-        <Messages messages={messages} status={status} />
+        <Messages
+          messages={messages}
+          status={status}
+          linkedAccountOwnerId={selectedLinkedAccountOwnerId}
+          apiKey={apiKey}
+          addToolResult={handleAddToolResult}
+        />
         <ChatInput
           input={input}
+          setMessages={setMessages}
+          stop={stop}
           handleInputChange={handleInputChange}
           handleSubmit={handleSubmit}
           status={status}