feat: Implement user confirmation mode, request confirmation when run…

…ning bash/python code in this mode (All-Hands-AI#2774) * [feat] confirmation mode for bash actions * feat: Add modal setting for Confirmation Mode * fix: frontend tests for confirmation mode switch * fix: add missing CONFIRMATION_MODE value in SettingsModal.test.tsx * fix: update test to integrate new setting * feat: Implement user confirmation for running bash/python code * fix: don't display rejected actions * fix: linting, rename/refactor based on feedback * fix: add property only to commands, pass serialization tests * fix: package-lock.json, lint test_action_serialization.py * test: add is_confirmed to integration test outputs --------- Co-authored-by: Mislav Balunovic <mislav.balunovic@gmail.com>
invariantlabs-ai · Jul 11, 2024 · 5f61885 · 5f61885
1 parent 1d4f422
commit 5f61885
Show file tree

Hide file tree

Showing 60 changed files with 524 additions and 106 deletions.
diff --git a/frontend/src/assets/confirm.tsx b/frontend/src/assets/confirm.tsx
@@ -0,0 +1,18 @@
+import React from "react";
+
+function ConfirmIcon(): JSX.Element {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      fill="none"
+      viewBox="0 0 24 24"
+      strokeWidth={1.5}
+      stroke="currentColor"
+      className="w-5 h-5"
+    >
+      <path strokeLinecap="round" strokeLinejoin="round" d="M5 13l4 4L19 7" />
+    </svg>
+  );
+}
+
+export default ConfirmIcon;
diff --git a/frontend/src/assets/reject.tsx b/frontend/src/assets/reject.tsx
@@ -0,0 +1,22 @@
+import React from "react";
+
+function RejectIcon(): JSX.Element {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      fill="none"
+      viewBox="0 0 24 24"
+      strokeWidth={1.5}
+      stroke="currentColor"
+      className="w-5 h-5"
+    >
+      <path
+        strokeLinecap="round"
+        strokeLinejoin="round"
+        d="M6 18L18 6M6 6l12 12"
+      />
+    </svg>
+  );
+}
+
+export default RejectIcon;
diff --git a/frontend/src/components/AgentControlBar.tsx b/frontend/src/components/AgentControlBar.tsx
@@ -17,6 +17,7 @@ const IgnoreTaskStateMap: { [k: string]: AgentState[] } = {
     AgentState.FINISHED,
     AgentState.REJECTED,
     AgentState.AWAITING_USER_INPUT,
+    AgentState.AWAITING_USER_CONFIRMATION,
   ],
   [AgentState.RUNNING]: [
     AgentState.INIT,
@@ -25,8 +26,12 @@ const IgnoreTaskStateMap: { [k: string]: AgentState[] } = {
     AgentState.FINISHED,
     AgentState.REJECTED,
     AgentState.AWAITING_USER_INPUT,
+    AgentState.AWAITING_USER_CONFIRMATION,
   ],
   [AgentState.STOPPED]: [AgentState.INIT, AgentState.STOPPED],
+  [AgentState.USER_CONFIRMED]: [AgentState.RUNNING],
+  [AgentState.USER_REJECTED]: [AgentState.RUNNING],
+  [AgentState.AWAITING_USER_CONFIRMATION]: [],
 };
 
 interface ButtonProps {
@@ -101,42 +106,44 @@ function AgentControlBar() {
   }, [curAgentState]);
 
   return (
-    <div className="flex items-center gap-3">
-      {curAgentState === AgentState.PAUSED ? (
+    <div className="flex justify-between items-center gap-20">
+      <div className="flex items-center gap-3">
+        {curAgentState === AgentState.PAUSED ? (
+          <ActionButton
+            isDisabled={
+              isLoading ||
+              IgnoreTaskStateMap[AgentState.RUNNING].includes(curAgentState)
+            }
+            content="Resume the agent task"
+            action={AgentState.RUNNING}
+            handleAction={handleAction}
+            large
+          >
+            <PlayIcon />
+          </ActionButton>
+        ) : (
+          <ActionButton
+            isDisabled={
+              isLoading ||
+              IgnoreTaskStateMap[AgentState.PAUSED].includes(curAgentState)
+            }
+            content="Pause the current task"
+            action={AgentState.PAUSED}
+            handleAction={handleAction}
+            large
+          >
+            <PauseIcon />
+          </ActionButton>
+        )}
         <ActionButton
-          isDisabled={
-            isLoading ||
-            IgnoreTaskStateMap[AgentState.RUNNING].includes(curAgentState)
-          }
-          content="Resume the agent task"
-          action={AgentState.RUNNING}
+          isDisabled={isLoading}
+          content="Start a new task"
+          action={AgentState.STOPPED}
           handleAction={handleAction}
-          large
         >
-          <PlayIcon />
+          <ArrowIcon />
         </ActionButton>
-      ) : (
-        <ActionButton
-          isDisabled={
-            isLoading ||
-            IgnoreTaskStateMap[AgentState.PAUSED].includes(curAgentState)
-          }
-          content="Pause the current task"
-          action={AgentState.PAUSED}
-          handleAction={handleAction}
-          large
-        >
-          <PauseIcon />
-        </ActionButton>
-      )}
-      <ActionButton
-        isDisabled={isLoading}
-        content="Start a new task"
-        action={AgentState.STOPPED}
-        handleAction={handleAction}
-      >
-        <ArrowIcon />
-      </ActionButton>
+      </div>
     </div>
   );
 }

diff --git a/frontend/src/components/AgentStatusBar.tsx b/frontend/src/components/AgentStatusBar.tsx
@@ -58,6 +58,20 @@ function AgentStatusBar() {
       message: t(I18nKey.CHAT_INTERFACE$AGENT_ERROR_MESSAGE),
       indicator: IndicatorColor.RED,
     },
+    [AgentState.AWAITING_USER_CONFIRMATION]: {
+      message: t(
+        I18nKey.CHAT_INTERFACE$AGENT_AWAITING_USER_CONFIRMATION_MESSAGE,
+      ),
+      indicator: IndicatorColor.ORANGE,
+    },
+    [AgentState.USER_CONFIRMED]: {
+      message: t(I18nKey.CHAT_INTERFACE$AGENT_ACTION_USER_CONFIRMED_MESSAGE),
+      indicator: IndicatorColor.GREEN,
+    },
+    [AgentState.USER_REJECTED]: {
+      message: t(I18nKey.CHAT_INTERFACE$AGENT_ACTION_USER_REJECTED_MESSAGE),
+      indicator: IndicatorColor.RED,
+    },
   };
 
   // TODO: Extend the agent status, e.g.:

diff --git a/frontend/src/components/chat/Chat.test.tsx b/frontend/src/components/chat/Chat.test.tsx
@@ -1,6 +1,7 @@
 import React from "react";
-import { render, screen } from "@testing-library/react";
+import { screen } from "@testing-library/react";
 import { describe, expect, it } from "vitest";
+import { renderWithProviders } from "test-utils";
 import Chat from "./Chat";
 
 const MESSAGES: Message[] = [
@@ -13,7 +14,7 @@ HTMLElement.prototype.scrollTo = vi.fn(() => {});
 
 describe("Chat", () => {
   it("should render chat messages", () => {
-    render(<Chat messages={MESSAGES} />);
+    renderWithProviders(<Chat messages={MESSAGES} />);
 
     const messages = screen.getAllByTestId("message");
 

diff --git a/frontend/src/components/chat/Chat.tsx b/frontend/src/components/chat/Chat.tsx
@@ -1,15 +1,24 @@
 import React from "react";
 import ChatMessage from "./ChatMessage";
+import AgentState from "#/types/AgentState";
 
 interface ChatProps {
   messages: Message[];
+  curAgentState?: AgentState;
 }
 
-function Chat({ messages }: ChatProps) {
+function Chat({ messages, curAgentState }: ChatProps) {
   return (
     <div className="flex flex-col gap-3">
       {messages.map((message, index) => (
-        <ChatMessage key={index} message={message} />
+        <ChatMessage
+          key={index}
+          message={message}
+          isLastMessage={messages && index === messages.length - 1}
+          awaitingUserConfirmation={
+            curAgentState === AgentState.AWAITING_USER_CONFIRMATION
+          }
+        />
       ))}
     </div>
   );

diff --git a/frontend/src/components/chat/ChatInterface.tsx b/frontend/src/components/chat/ChatInterface.tsx
@@ -123,7 +123,7 @@ function ChatInterface() {
           className="overflow-y-auto p-3"
           onScroll={(e) => onChatBodyScroll(e.currentTarget)}
         >
-          <Chat messages={messages} />
+          <Chat messages={messages} curAgentState={curAgentState} />
         </div>
       </div>
 
@@ -169,7 +169,10 @@ function ChatInterface() {
       </div>
 
       <ChatInput
-        disabled={curAgentState === AgentState.LOADING}
+        disabled={
+          curAgentState === AgentState.LOADING ||
+          curAgentState === AgentState.AWAITING_USER_CONFIRMATION
+        }
         onSendMessage={handleSendMessage}
       />
       <FeedbackModal

diff --git a/frontend/src/components/chat/ChatMessage.test.tsx b/frontend/src/components/chat/ChatMessage.test.tsx
@@ -10,14 +10,24 @@ vi.mock("#/hooks/useTyping", () => ({
 
 describe("Message", () => {
   it("should render a user message", () => {
-    render(<ChatMessage message={{ sender: "user", content: "Hello" }} />);
+    render(
+      <ChatMessage
+        message={{ sender: "user", content: "Hello" }}
+        isLastMessage={false}
+      />,
+    );
 
     expect(screen.getByTestId("message")).toBeInTheDocument();
     expect(screen.getByTestId("message")).toHaveClass("self-end"); // user message should be on the right side
   });
 
   it("should render an assistant message", () => {
-    render(<ChatMessage message={{ sender: "assistant", content: "Hi" }} />);
+    render(
+      <ChatMessage
+        message={{ sender: "assistant", content: "Hi" }}
+        isLastMessage={false}
+      />,
+    );
 
     expect(screen.getByTestId("message")).toBeInTheDocument();
     expect(screen.getByTestId("message")).not.toHaveClass("self-end"); // assistant message should be on the left side
@@ -30,6 +40,7 @@ describe("Message", () => {
           sender: "user",
           content: "```js\nconsole.log('Hello')\n```",
         }}
+        isLastMessage={false}
       />,
     );
 

diff --git a/frontend/src/components/chat/ChatMessage.tsx b/frontend/src/components/chat/ChatMessage.tsx
@@ -3,15 +3,26 @@ import Markdown from "react-markdown";
 import { FaClipboard, FaClipboardCheck } from "react-icons/fa";
 import { twMerge } from "tailwind-merge";
 import { useTranslation } from "react-i18next";
+import { Tooltip } from "@nextui-org/react";
+import AgentState from "#/types/AgentState";
 import { code } from "../markdown/code";
 import toast from "#/utils/toast";
 import { I18nKey } from "#/i18n/declaration";
+import ConfirmIcon from "#/assets/confirm";
+import RejectIcon from "#/assets/reject";
+import { changeAgentState } from "#/services/agentStateService";
 
 interface MessageProps {
   message: Message;
+  isLastMessage?: boolean;
+  awaitingUserConfirmation?: boolean;
 }
 
-function ChatMessage({ message }: MessageProps) {
+function ChatMessage({
+  message,
+  isLastMessage,
+  awaitingUserConfirmation,
+}: MessageProps) {
   const [isCopy, setIsCopy] = useState(false);
   const [isHovering, setIsHovering] = useState(false);
 
@@ -58,6 +69,45 @@ function ChatMessage({ message }: MessageProps) {
         </button>
       )}
       <Markdown components={{ code }}>{message.content}</Markdown>
+      {isLastMessage &&
+        message.sender === "assistant" &&
+        awaitingUserConfirmation && (
+          <div className="flex justify-between items-center pt-4">
+            <p>{t(I18nKey.CHAT_INTERFACE$USER_ASK_CONFIRMATION)}</p>
+            <div className="flex items-center gap-3">
+              <Tooltip
+                content={t(I18nKey.CHAT_INTERFACE$USER_CONFIRMED)}
+                closeDelay={100}
+              >
+                <button
+                  type="button"
+                  aria-label="Confirm action"
+                  className="bg-neutral-700 rounded-full p-1 hover:bg-neutral-800"
+                  onClick={() => {
+                    changeAgentState(AgentState.USER_CONFIRMED);
+                  }}
+                >
+                  <ConfirmIcon />
+                </button>
+              </Tooltip>
+              <Tooltip
+                content={t(I18nKey.CHAT_INTERFACE$USER_REJECTED)}
+                closeDelay={100}
+              >
+                <button
+                  type="button"
+                  aria-label="Reject action"
+                  className="bg-neutral-700 rounded-full p-1 hover:bg-neutral-800"
+                  onClick={() => {
+                    changeAgentState(AgentState.USER_REJECTED);
+                  }}
+                >
+                  <RejectIcon />
+                </button>
+              </Tooltip>
+            </div>
+          </div>
+        )}
     </div>
   );
 }