Make async

moatless-ui/src/pages/runs/[id].tsx

@@ -7,29 +7,27 @@ import { Alert, AlertDescription } from '@/lib/components/ui/alert';
 import { ResizableHandle, ResizablePanel, ResizablePanelGroup } from '@/lib/components/ui/resizable';
 import { RunStatus } from './components/RunStatus';
 import { RunEvents } from './components/RunEvents';
-import { useWebSocketStore } from '@/lib/stores/websocketStore';
-import { useMemo } from 'react';
+import { useEffect } from 'react';
 import { ScrollArea } from '@/lib/components/ui/scroll-area';
 import { TimelineItemDetails } from './components/TimelineItemDetails';
+import { useQueryClient } from '@tanstack/react-query';
+import { useWebSocketStore } from '@/lib/stores/websocketStore';
 
 export function RunPage() {
   const { id } = useParams<{ id: string }>();
   const { data: runData, isError, error } = useRun(id!);
+  const queryClient = useQueryClient();
+  const { subscribe } = useWebSocketStore();
 
-  // Use useMemo to cache the selector functions
-  //const selectMessages = useMemo(
-  //  () => (state: any) => state.messages,
-  //  [id]
-  //);
-
-  const selectConnectionStatus = useMemo(
-    () => (state: any) => state.connectionStatus,
-    []
-  );
+  useEffect(() => {
+    if (!id) return;
+
+    const unsubscribe = subscribe(`run.${id}`, () => {
+      queryClient.invalidateQueries({ queryKey: ['run', id] });
+    });
 
-  // Use the memoized selectors
-  //const messages = useWebSocketStore(selectMessages);
-  const wsStatus = useWebSocketStore(selectConnectionStatus);
+    return () => unsubscribe();
+  }, [id, subscribe, queryClient]);
 
   if (isError) {
     return (
@@ -59,16 +57,6 @@ export function RunPage() {
     );
   }
 
-  // Combine WebSocket messages with initial events
-  //const allEvents = useMemo(() => {
-  //  const wsEvents = messages.map(msg => ({
-  //    event_type: msg.type,
-  //    timestamp: new Date().toISOString(),
-  //    data: { message: msg.message || msg.error }
-  //    }));
-  //  return [...(runData.events || []), ...wsEvents];
-  //}, [runData.events, messages]);
-
   return (
     <div className="h-[calc(100vh-56px)]">
       <ResizablePanelGroup 

moatless-ui/src/pages/runs/components/RunEvents.tsx

@@ -1,51 +1,78 @@
 import { RunEvent } from '@/lib/types/run';
-import { AlertCircle, Info, MessageSquare, Bot, Terminal } from 'lucide-react';
+import { AlertCircle, Info, MessageSquare, Bot, Terminal, Loader2 } from 'lucide-react';
 import { formatDistanceToNow } from 'date-fns';
+import { ScrollArea } from '@/lib/components/ui/scroll-area';
+import { cn } from '@/lib/utils';
 
 interface RunEventsProps {
   events: RunEvent[];
+  className?: string;
 }
 
-export function RunEvents({ events }: RunEventsProps) {
-  // desc sort events by timestamp
-  const sortedEvents = events.sort((a, b) => b.timestamp - a.timestamp);
+export function RunEvents({ events, className }: RunEventsProps) {
+  // just sort in reverse order
+  const reversedEvents = [...events].reverse();
 
   const getEventIcon = (type: string) => {
     switch (type) {
       case 'error':
         return <AlertCircle className="h-4 w-4 text-destructive" />;
       case 'agent_message':
-        return <Bot className="h-4 w-4" />;
+        return <Bot className="h-4 w-4 text-primary" />;
       case 'system_message':
-        return <Terminal className="h-4 w-4" />;
+        return <Terminal className="h-4 w-4 text-muted-foreground" />;
       case 'user_message':
-        return <MessageSquare className="h-4 w-4" />;
+        return <MessageSquare className="h-4 w-4 text-blue-500" />;
+      case 'agent_action_created':
+        return <Bot className="h-4 w-4 text-primary" />;
+      case 'agent_action_executed':
+        return <Loader2 className="h-4 w-4 text-green-500" />;
       default:
-        return <Info className="h-4 w-4" />;
+        return <Info className="h-4 w-4 text-muted-foreground" />;
     }
   };
 
+  const formatEventType = (type: string) => {
+    return type
+      .split('_')
+      .map(word => word.charAt(0).toUpperCase() + word.slice(1))
+      .join(' ');
+  };
+
+  if (events.length === 0) {
+    return (
+      <div className="flex items-center justify-center h-32 text-muted-foreground">
+        No events yet
+      </div>
+    );
+  }
+
   return (
     <div className="space-y-2 p-4">
-      {sortedEvents.map((event, i) => (
+      {reversedEvents.map((event, i) => (
         <div 
-          key={i} 
-          className="rounded-md bg-muted p-3 text-sm"
+          key={`${event.timestamp}-${i}`}
+          className="rounded-md border bg-card p-3 text-sm shadow-sm"
         >
           <div className="flex items-center justify-between mb-1">
             <div className="flex items-center gap-2">
               {getEventIcon(event.event_type)}
               <span className="font-medium">
-                {event.event_type.replace(/_/g, ' ')}
+                {formatEventType(event.event_type)}
               </span>
+              {event.node_id !== undefined && (
+                <span className="text-xs text-muted-foreground">
+                  Node {event.node_id}
+                </span>
+              )}
             </div>
             <span className="text-xs text-muted-foreground">
               {formatDistanceToNow(new Date(event.timestamp))} ago
             </span>
           </div>
-          {event.data.message && (
-            <p className="text-sm text-muted-foreground mt-1">
-              {event.data.message}
+          {event.action_name && (
+            <p className="text-xs text-muted-foreground mt-1">
+              Action: {event.action_name}
             </p>
           )}
         </div>

moatless-ui/src/pages/validate/index.tsx

@@ -1,4 +1,4 @@
-import { useState } from 'react';
+import { useState, useEffect } from 'react';
 import { Button } from '@/lib/components/ui/button';
 import { Loader2 } from 'lucide-react';
 import { toast } from 'sonner';
@@ -10,6 +10,8 @@ import { ModelSelector } from '@/lib/components/selectors/ModelSelector';
 import { useValidationStore } from '@/stores/validationStore';
 import { AlertCircle } from 'lucide-react';
 import { Alert, AlertDescription, AlertTitle } from '@/lib/components/ui/alert';
+import { useWebSocket } from '@/lib/stores/websocketStore';
+import { useQueryClient } from '@tanstack/react-query';
 
 export function ValidatePage() {
   const navigate = useNavigate();
@@ -31,9 +33,12 @@ export function ValidatePage() {
   // Add error state
   const [error, setError] = useState<string | null>(null);
 
+  const queryClient = useQueryClient();
+  const { subscribe, unsubscribe } = useWebSocket();
+
   const handleSubmit = async (e: React.FormEvent) => {
     e.preventDefault();
-    setError(null); // Clear any previous errors
+    setError(null);
 
     try {
       const response = await startValidation.mutateAsync({
@@ -56,14 +61,6 @@ export function ValidatePage() {
     }
   };
 
-  if (startValidation.isLoading) {
-    return (
-      <div className="flex h-full items-center justify-center">
-        <Loader2 className="h-6 w-6 animate-spin" />
-      </div>
-    );
-  }
-
   return (
     <div className="container mx-auto py-6">
       <h1 className="mb-6 text-2xl font-bold">Validate Agent</h1>
@@ -97,9 +94,16 @@ export function ValidatePage() {
         <div className="flex justify-end">
           <Button
             type="submit"
-            disabled={!selectedAgentId || !selectedModelId || !selectedInstanceId}
+            disabled={!selectedAgentId || !selectedModelId || !selectedInstanceId || startValidation.isPending}
           >
-            Run Validation
+            {startValidation.isPending ? (
+              <>
+                <Loader2 className="mr-2 h-4 w-4 animate-spin" />
+                Starting...
+              </>
+            ) : (
+              'Run Validation'
+            )}
           </Button>
         </div>
       </form>

moatless/actions/action.py

@@ -56,7 +56,7 @@ class Action(BaseModel, ABC):
 
     _workspace: Workspace = PrivateAttr(default=None)
 
-    def execute(self, args: ActionArguments, file_context: FileContext | None = None) -> Observation:
+    async def execute(self, args: ActionArguments, file_context: FileContext | None = None) -> Observation:
         """
         Execute the action.
         """
@@ -67,7 +67,7 @@ def execute(self, args: ActionArguments, file_context: FileContext | None = None
         message = self._execute(args, file_context=file_context)
         return Observation.create(message)
 
-    def _execute(self, args: ActionArguments, file_context: FileContext | None = None) -> str | None:
+    async def _execute(self, args: ActionArguments, file_context: FileContext | None = None) -> str | None:
         """
         Execute the action and return the updated FileContext.
         """

moatless/actions/append_string.py

@@ -62,7 +62,7 @@ class AppendString(Action, CodeActionValueMixin, CodeModificationMixin):
 
     args_schema = AppendStringArgs
 
-    def execute(
+    async def execute(
         self,
         args: AppendStringArgs,
         file_context: FileContext | None = None,

moatless/actions/claude_text_editor.py

@@ -153,7 +153,7 @@ class ClaudeEditTool(Action, CodeModificationMixin):
     _create_file: CreateFile = PrivateAttr()
     _repository: Repository | None = PrivateAttr(None)
 
-    def execute(
+    async def execute(
         self,
         args: EditActionArguments,
         file_context: FileContext | None = None,

moatless/actions/create_file.py

@@ -51,7 +51,7 @@ class CreateFile(Action, CodeActionValueMixin, CodeModificationMixin):
 
     args_schema = CreateFileArgs
 
-    def execute(
+    async def execute(
         self,
         args: CreateFileArgs,
         file_context: FileContext | None = None,

moatless/actions/finish.py

@@ -46,7 +46,7 @@ class Finish(Action):
         description="Whether to enforce that the file context has a test patch",
     )
 
-    def execute(
+    async def execute(
         self,
         args: FinishArgs,
         file_context: FileContext | None = None,

moatless/actions/identify_mixin.py

@@ -73,7 +73,7 @@ class IdentifyMixin(CompletionModelMixin):
         description="The maximum number of tokens allowed in the identify prompt.",
     )
 
-    def _identify_code(self, args, view_context: FileContext, max_tokens: int) -> Tuple[FileContext, Completion]:
+    async def _identify_code(self, args, view_context: FileContext, max_tokens: int) -> Tuple[FileContext, Completion]:
         """Identify relevant code sections in a large context.
 
         Args:
@@ -106,7 +106,7 @@ def _identify_code(self, args, view_context: FileContext, max_tokens: int) -> Tu
 
         MAX_RETRIES = 3
         for retry_attempt in range(MAX_RETRIES):
-            completion_response = self._completion_model.create_completion(messages=messages)
+            completion_response = await self._completion_model.create_completion(messages=messages)
             logger.info(
                 f"Identifying relevant code sections. Attempt {retry_attempt + 1} of {MAX_RETRIES}.{len(completion_response.structured_outputs)} identify requests."
             )

moatless/actions/insert_line.py

@@ -55,7 +55,7 @@ class InsertLine(Action, CodeActionValueMixin, CodeModificationMixin):
 
     args_schema = InsertLinesArgs
 
-    def execute(
+    async def execute(
         self,
         args: InsertLinesArgs,
         file_context: FileContext | None = None,

moatless/actions/list_files.py

@@ -33,7 +33,7 @@ def short_summary(self) -> str:
 class ListFiles(Action):
     args_schema = ListFilesArgs
 
-    def execute(
+    async def execute(
         self,
         args: ListFilesArgs,
         file_context: FileContext | None = None,

moatless/actions/reject.py

@@ -25,5 +25,5 @@ def equals(self, other: "ActionArguments") -> bool:
 class Reject(Action):
     args_schema: ClassVar[Type[ActionArguments]] = RejectArgs
 
-    def execute(self, args: RejectArgs, file_context: FileContext | None = None):
+    async def execute(self, args: RejectArgs, file_context: FileContext | None = None):
         return Observation(message=args.rejection_reason, terminal=True)

moatless/actions/respond.py

@@ -27,7 +27,7 @@ class MessageAction(Action):
 
     args_schema = MessageArgs
 
-    def execute(
+    async def execute(
         self,
         args: MessageArgs,
         file_context: FileContext | None = None,

moatless/actions/run_tests.py

@@ -47,7 +47,7 @@ class RunTests(Action):
     _repository: Repository = PrivateAttr()
     _runtime: RuntimeEnvironment = PrivateAttr()
 
-    def execute(
+    async def execute(
         self,
         args: RunTestsArgs,
         file_context: FileContext | None = None,

moatless/actions/search_base.py

@@ -106,7 +106,7 @@ def _initialize_completion_model(self):
         if self._completion_model:
             self._completion_model.initialize(Identify, IDENTIFY_SYSTEM_PROMPT)
 
-    def execute(self, args: SearchBaseArgs, file_context: FileContext | None = None) -> Observation:
+    async def execute(self, args: SearchBaseArgs, file_context: FileContext | None = None) -> Observation:
         if file_context is None:
             raise ValueError("File context must be provided to execute the search action.")
 
@@ -141,7 +141,7 @@ def execute(self, args: SearchBaseArgs, file_context: FileContext | None = None)
             logger.info(
                 f"{self.name}: Search too large. {properties['search_tokens']} tokens and {search_result_context.span_count()} hits, will ask for clarification."
             )
-            view_context, completion = self._identify_code(args, search_result_context)
+            view_context, completion = await self._identify_code(args, search_result_context)
         else:
             view_context = search_result_context
 
@@ -239,7 +239,7 @@ def _search(self, args: SearchBaseArgs) -> SearchCodeResponse:
     def _search_for_alternative_suggestion(self, args: SearchBaseArgs) -> SearchCodeResponse:
         return SearchCodeResponse()
 
-    def _identify_code(
+    async def _identify_code(
         self, args: SearchBaseArgs, search_result_ctx: FileContext
     ) -> Tuple[IdentifiedSpans, Completion]:
         search_result_str = search_result_ctx.create_prompt(
@@ -263,7 +263,7 @@ def _identify_code(
 
         MAX_RETRIES = 3
         for retry_attempt in range(MAX_RETRIES):
-            completion_response = self.completion_model.create_completion(messages=messages)
+            completion_response = await self.completion_model.create_completion(messages=messages)
             logger.info(
                 f"Identifying relevant code sections. Attempt {retry_attempt + 1} of {MAX_RETRIES}.{len(completion_response.structured_outputs)} identify requests."
             )

moatless/actions/string_replace.py

@@ -132,7 +132,7 @@ class StringReplace(Action, CodeActionValueMixin, CodeModificationMixin):
         description="When True, automatically corrects indentation if all lines have the same indentation difference",
     )
 
-    def execute(self, args: StringReplaceArgs, file_context: FileContext | None = None) -> Observation:
+    async def execute(self, args: StringReplaceArgs, file_context: FileContext | None = None) -> Observation:
         path_str = self.normalize_path(args.path)
         path, error = self.validate_file_access(path_str, file_context)
         if error:

moatless/actions/verified_finish.py

@@ -43,7 +43,7 @@ def equals(self, other: "ActionArguments") -> bool:
 class VerifiedFinish(Action):
     args_schema: ClassVar[Type[ActionArguments]] = VerifiedFinishArgs
 
-    def execute(
+    async def execute(
         self,
         args: VerifiedFinishArgs,
         file_context: FileContext | None = None,