[Assist] Fix command output showing when running on multiple nodes

web/packages/teleport/src/Assist/context/AssistContext.tsx

@@ -409,15 +409,7 @@ export function AssistContextProvider(props: PropsWithChildren<unknown>) {
     executeCommandWebSocket.current = new WebSocket(url);
     executeCommandWebSocket.current.binaryType = 'arraybuffer';
 
-    executeCommandWebSocket.current.onclose = () => {
-      for (const nodeId of nodeIdToResultId.keys()) {
-        dispatch({
-          type: AssistStateActionType.FinishCommandResult,
-          conversationId: state.conversations.selectedId,
-          commandResultId: nodeIdToResultId.get(nodeId),
-        });
-      }
-    };
+    let sessionsEnded = 0;
 
     executeCommandWebSocket.current.onmessage = event => {
       const uintArray = new Uint8Array(event.data);
@@ -457,16 +449,26 @@ export function AssistContextProvider(props: PropsWithChildren<unknown>) {
           break;
 
         case MessageTypeEnum.SESSION_END:
-          for (const nodeId of nodeIdToResultId.keys()) {
-            dispatch({
-              type: AssistStateActionType.FinishCommandResult,
-              conversationId: state.conversations.selectedId,
-              commandResultId: nodeIdToResultId.get(nodeId),
-            });
+          // we don't know the nodeId of the session that ended, so we have to
+          // count the finished sessions and then mark them all as done once
+          // they've all finished
+          sessionsEnded += 1;
+
+          if (sessionsEnded === nodeIdToResultId.size) {
+            for (const nodeId of nodeIdToResultId.keys()) {
+              dispatch({
+                type: AssistStateActionType.FinishCommandResult,
+                conversationId: state.conversations.selectedId,
+                commandResultId: nodeIdToResultId.get(nodeId),
+              });
+            }
+
+            nodeIdToResultId.clear();
+
+            // TODO(ryan): move this to after the summary is sent once it's implemented
+            executeCommandWebSocket.current.close();
           }
 
-          nodeIdToResultId.clear();
-
           break;
       }
     };