Separate agent controller and server via EventStream (#1538)

* move towards event stream

* refactor agent state changes

* move agent state logic

* fix callbacks

* break on finish

* closer to working

* change frontend to accomodate new flow

* handle start action

* fix locked stream

* revert message

* logspam

* no async on close

* get rid of agent_task

* fix up closing

* better asyncio handling

* sleep to give back control

* fix key

* logspam

* update frontend agent state actions

* fix pause and cancel

* delint

* fix map

* delint

* wait for agent to finish

* fix unit test

* event stream enums

* fix merge issues

* fix lint

* fix test

* fix test

* add user message action

* add user message action

* fix up user messages

* fix main.py flow

* refactor message waiting

* lint

* fix test

* fix test

Author: rbren

frontend/src/components/AgentControlBar.tsx

@@ -4,47 +4,39 @@ import { useSelector } from "react-redux";
 import ArrowIcon from "#/assets/arrow";
 import PauseIcon from "#/assets/pause";
 import PlayIcon from "#/assets/play";
-import { changeTaskState } from "#/services/agentStateService";
+import { changeAgentState } from "#/services/agentStateService";
 import { clearMsgs } from "#/services/session";
 import store, { RootState } from "#/store";
-import AgentTaskAction from "#/types/AgentTaskAction";
-import AgentTaskState from "#/types/AgentTaskState";
+import AgentState from "#/types/AgentState";
 import { clearMessages } from "#/state/chatSlice";
 
-const TaskStateActionMap = {
-  [AgentTaskAction.START]: AgentTaskState.RUNNING,
-  [AgentTaskAction.PAUSE]: AgentTaskState.PAUSED,
-  [AgentTaskAction.RESUME]: AgentTaskState.RUNNING,
-  [AgentTaskAction.STOP]: AgentTaskState.STOPPED,
-};
-
-const IgnoreTaskStateMap: { [k: string]: AgentTaskState[] } = {
-  [AgentTaskAction.PAUSE]: [
-    AgentTaskState.INIT,
-    AgentTaskState.PAUSED,
-    AgentTaskState.STOPPED,
-    AgentTaskState.FINISHED,
-    AgentTaskState.AWAITING_USER_INPUT,
+const IgnoreTaskStateMap: { [k: string]: AgentState[] } = {
+  [AgentState.PAUSED]: [
+    AgentState.INIT,
+    AgentState.PAUSED,
+    AgentState.STOPPED,
+    AgentState.FINISHED,
+    AgentState.AWAITING_USER_INPUT,
   ],
-  [AgentTaskAction.RESUME]: [
-    AgentTaskState.INIT,
-    AgentTaskState.RUNNING,
-    AgentTaskState.STOPPED,
-    AgentTaskState.FINISHED,
-    AgentTaskState.AWAITING_USER_INPUT,
+  [AgentState.RUNNING]: [
+    AgentState.INIT,
+    AgentState.RUNNING,
+    AgentState.STOPPED,
+    AgentState.FINISHED,
+    AgentState.AWAITING_USER_INPUT,
   ],
-  [AgentTaskAction.STOP]: [
-    AgentTaskState.INIT,
-    AgentTaskState.STOPPED,
-    AgentTaskState.FINISHED,
+  [AgentState.STOPPED]: [
+    AgentState.INIT,
+    AgentState.STOPPED,
+    AgentState.FINISHED,
   ],
 };
 
 interface ButtonProps {
   isDisabled: boolean;
   content: string;
-  action: AgentTaskAction;
-  handleAction: (action: AgentTaskAction) => void;
+  action: AgentState;
+  handleAction: (action: AgentState) => void;
   large?: boolean;
 }
 
@@ -75,53 +67,53 @@ ActionButton.defaultProps = {
 };
 
 function AgentControlBar() {
-  const { curTaskState } = useSelector((state: RootState) => state.agent);
-  const [desiredState, setDesiredState] = React.useState(AgentTaskState.INIT);
+  const { curAgentState } = useSelector((state: RootState) => state.agent);
+  const [desiredState, setDesiredState] = React.useState(AgentState.INIT);
   const [isLoading, setIsLoading] = React.useState(false);
 
-  const handleAction = (action: AgentTaskAction) => {
-    if (IgnoreTaskStateMap[action].includes(curTaskState)) {
+  const handleAction = (action: AgentState) => {
+    if (IgnoreTaskStateMap[action].includes(curAgentState)) {
       return;
     }
 
     let act = action;
 
-    if (act === AgentTaskAction.STOP) {
-      act = AgentTaskAction.STOP;
+    if (act === AgentState.STOPPED) {
+      act = AgentState.STOPPED;
       clearMsgs().then().catch();
       store.dispatch(clearMessages());
     } else {
       setIsLoading(true);
     }
 
-    setDesiredState(TaskStateActionMap[act]);
-    changeTaskState(act);
+    setDesiredState(act);
+    changeAgentState(act);
   };
 
   useEffect(() => {
-    if (curTaskState === desiredState) {
-      if (curTaskState === AgentTaskState.STOPPED) {
+    if (curAgentState === desiredState) {
+      if (curAgentState === AgentState.STOPPED) {
         clearMsgs().then().catch();
         store.dispatch(clearMessages());
       }
       setIsLoading(false);
-    } else if (curTaskState === AgentTaskState.RUNNING) {
-      setDesiredState(AgentTaskState.RUNNING);
+    } else if (curAgentState === AgentState.RUNNING) {
+      setDesiredState(AgentState.RUNNING);
     }
-    // We only want to run this effect when curTaskState changes
+    // We only want to run this effect when curAgentState changes
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [curTaskState]);
+  }, [curAgentState]);
 
   return (
     <div className="flex items-center gap-3">
-      {curTaskState === AgentTaskState.PAUSED ? (
+      {curAgentState === AgentState.PAUSED ? (
         <ActionButton
           isDisabled={
             isLoading ||
-            IgnoreTaskStateMap[AgentTaskAction.RESUME].includes(curTaskState)
+            IgnoreTaskStateMap[AgentState.RUNNING].includes(curAgentState)
           }
           content="Resume the agent task"
-          action={AgentTaskAction.RESUME}
+          action={AgentState.RUNNING}
           handleAction={handleAction}
           large
         >
@@ -131,10 +123,10 @@ function AgentControlBar() {
         <ActionButton
           isDisabled={
             isLoading ||
-            IgnoreTaskStateMap[AgentTaskAction.PAUSE].includes(curTaskState)
+            IgnoreTaskStateMap[AgentState.PAUSED].includes(curAgentState)
           }
           content="Pause the agent task"
-          action={AgentTaskAction.PAUSE}
+          action={AgentState.PAUSED}
           handleAction={handleAction}
           large
         >
@@ -144,7 +136,7 @@ function AgentControlBar() {
       <ActionButton
         isDisabled={isLoading}
         content="Restart a new agent task"
-        action={AgentTaskAction.STOP}
+        action={AgentState.STOPPED}
         handleAction={handleAction}
       >
         <ArrowIcon />

frontend/src/components/AgentStatusBar.tsx

@@ -3,40 +3,39 @@ import { useTranslation } from "react-i18next";
 import { useSelector } from "react-redux";
 import { I18nKey } from "#/i18n/declaration";
 import { RootState } from "#/store";
-import AgentTaskState from "#/types/AgentTaskState";
+import AgentState from "#/types/AgentState";
 
 const AgentStatusMap: { [k: string]: { message: string; indicator: string } } =
   {
-    [AgentTaskState.INIT]: {
+    [AgentState.INIT]: {
       message: "Agent is initialized, waiting for task...",
       indicator: "bg-blue-500",
     },
-    [AgentTaskState.RUNNING]: {
+    [AgentState.RUNNING]: {
       message: "Agent is running task...",
       indicator: "bg-green-500",
     },
-    [AgentTaskState.AWAITING_USER_INPUT]: {
+    [AgentState.AWAITING_USER_INPUT]: {
       message: "Agent is awaiting user input...",
       indicator: "bg-orange-500",
     },
-    [AgentTaskState.PAUSED]: {
+    [AgentState.PAUSED]: {
       message: "Agent has paused.",
       indicator: "bg-yellow-500",
     },
-    [AgentTaskState.STOPPED]: {
+    [AgentState.STOPPED]: {
       message: "Agent has stopped.",
       indicator: "bg-red-500",
     },
-    [AgentTaskState.FINISHED]: {
+    [AgentState.FINISHED]: {
       message: "Agent has finished the task.",
       indicator: "bg-green-500",
     },
   };
 
 function AgentStatusBar() {
   const { t } = useTranslation();
-  const { initialized } = useSelector((state: RootState) => state.task);
-  const { curTaskState } = useSelector((state: RootState) => state.agent);
+  const { curAgentState } = useSelector((state: RootState) => state.agent);
 
   // TODO: Extend the agent status, e.g.:
   // - Agent is typing
@@ -46,13 +45,13 @@ function AgentStatusBar() {
   // - Agent is not available
   return (
     <div className="flex items-center">
-      {initialized ? (
+      {curAgentState !== AgentState.LOADING ? (
         <>
           <div
-            className={`w-3 h-3 mr-2 rounded-full animate-pulse ${AgentStatusMap[curTaskState].indicator}`}
+            className={`w-3 h-3 mr-2 rounded-full animate-pulse ${AgentStatusMap[curAgentState].indicator}`}
           />
           <span className="text-sm text-stone-400">
-            {AgentStatusMap[curTaskState].message}
+            {AgentStatusMap[curAgentState].message}
           </span>
         </>
       ) : (

frontend/src/components/chat/ChatInterface.test.tsx

@@ -8,7 +8,7 @@ import ChatInterface from "./ChatInterface";
 import Socket from "#/services/socket";
 import ActionType from "#/types/ActionType";
 import { addAssistantMessage } from "#/state/chatSlice";
-import AgentTaskState from "#/types/AgentTaskState";
+import AgentState from "#/types/AgentState";
 
 // avoid typing side-effect
 vi.mock("#/hooks/useTyping", () => ({
@@ -25,7 +25,6 @@ const renderChatInterface = () =>
   renderWithProviders(<ChatInterface />, {
     preloadedState: {
       task: {
-        initialized: true,
         completed: false,
       },
     },
@@ -38,7 +37,16 @@ describe("ChatInterface", () => {
   });
 
   it("should render the new message the user has typed", async () => {
-    renderChatInterface();
+    renderWithProviders(<ChatInterface />, {
+      preloadedState: {
+        task: {
+          completed: false,
+        },
+        agent: {
+          curAgentState: AgentState.INIT,
+        },
+      },
+    });
 
     const input = screen.getByRole("textbox");
 
@@ -73,11 +81,10 @@ describe("ChatInterface", () => {
     renderWithProviders(<ChatInterface />, {
       preloadedState: {
         task: {
-          initialized: true,
           completed: false,
         },
         agent: {
-          curTaskState: AgentTaskState.INIT,
+          curAgentState: AgentState.INIT,
         },
       },
     });
@@ -95,11 +102,10 @@ describe("ChatInterface", () => {
     renderWithProviders(<ChatInterface />, {
       preloadedState: {
         task: {
-          initialized: true,
           completed: false,
         },
         agent: {
-          curTaskState: AgentTaskState.AWAITING_USER_INPUT,
+          curAgentState: AgentState.AWAITING_USER_INPUT,
         },
       },
     });
@@ -110,8 +116,8 @@ describe("ChatInterface", () => {
     });
 
     const event = {
-      action: ActionType.USER_MESSAGE,
-      args: { message: "my message" },
+      action: ActionType.MESSAGE,
+      args: { content: "my message" },
     };
     expect(socketSpy).toHaveBeenCalledWith(JSON.stringify(event));
   });
@@ -120,9 +126,11 @@ describe("ChatInterface", () => {
     renderWithProviders(<ChatInterface />, {
       preloadedState: {
         task: {
-          initialized: false,
           completed: false,
         },
+        agent: {
+          curAgentState: AgentState.LOADING,
+        },
       },
     });
 

frontend/src/components/chat/ChatInterface.tsx

@@ -1,32 +1,19 @@
 import React from "react";
-import { useDispatch, useSelector } from "react-redux";
+import { useSelector } from "react-redux";
 import { IoMdChatbubbles } from "react-icons/io";
 import ChatInput from "./ChatInput";
 import Chat from "./Chat";
 import { RootState } from "#/store";
-import AgentTaskState from "#/types/AgentTaskState";
-import { addUserMessage } from "#/state/chatSlice";
-import ActionType from "#/types/ActionType";
-import Socket from "#/services/socket";
+import AgentState from "#/types/AgentState";
+import { sendChatMessage } from "#/services/chatService";
 
 function ChatInterface() {
-  const { initialized } = useSelector((state: RootState) => state.task);
   const { messages } = useSelector((state: RootState) => state.chat);
-  const { curTaskState } = useSelector((state: RootState) => state.agent);
-
-  const dispatch = useDispatch();
+  const { curAgentState } = useSelector((state: RootState) => state.agent);
 
   const handleSendMessage = (content: string) => {
-    dispatch(addUserMessage(content));
-
-    let event;
-    if (curTaskState === AgentTaskState.INIT) {
-      event = { action: ActionType.START, args: { task: content } };
-    } else {
-      event = { action: ActionType.USER_MESSAGE, args: { message: content } };
-    }
-
-    Socket.send(JSON.stringify(event));
+    const isTask = curAgentState === AgentState.INIT;
+    sendChatMessage(content, isTask);
   };
 
   return (
@@ -42,7 +29,10 @@ function ChatInterface() {
         {/* Fade between messages and input */}
         <div className="absolute bottom-0 left-0 right-0 h-4 bg-gradient-to-b from-transparent to-neutral-800" />
       </div>
-      <ChatInput disabled={!initialized} onSendMessage={handleSendMessage} />
+      <ChatInput
+        disabled={curAgentState === AgentState.LOADING}
+        onSendMessage={handleSendMessage}
+      />
     </div>
   );
 }

frontend/src/components/modals/settings/SettingsForm.test.tsx

@@ -2,7 +2,7 @@ import { act, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import React from "react";
 import { renderWithProviders } from "test-utils";
-import AgentTaskState from "#/types/AgentTaskState";
+import AgentState from "#/types/AgentState";
 import { Settings } from "#/services/settings";
 import SettingsForm from "./SettingsForm";
 
@@ -80,7 +80,7 @@ describe("SettingsForm", () => {
         onLanguageChange={onLanguageChangeMock}
         onAPIKeyChange={onAPIKeyChangeMock}
       />,
-      { preloadedState: { agent: { curTaskState: AgentTaskState.RUNNING } } },
+      { preloadedState: { agent: { curAgentState: AgentState.RUNNING } } },
     );
     const modelInput = screen.getByRole("combobox", { name: "model" });
     const agentInput = screen.getByRole("combobox", { name: "agent" });