[Assist] Do not close the WS after command execution

lib/web/command.go

@@ -284,22 +284,9 @@ func (h *Handler) executeCommand(
 		return trace.Wrap(err)
 	}
 
-	// We want to send a session end event once, when nothing else will come
-	// through the websocket. We wait for all command exec to finish and the
-	// optional summary. Once the front-end gets this event it might also
-	// initiate a websocket close.
-	defer func() {
-		envelope := &Envelope{
-			Version: defaults.WebsocketVersion,
-			Type:    defaults.WebsocketClose,
-		}
-		envelopeBytes, _ := proto.Marshal(envelope)
-		_ = ws.WriteMessage(websocket.BinaryMessage, envelopeBytes)
-	}()
-
 	runCommands(hosts, runCmd, h.log)
 
-	// Optionally try to compute the command summary.
+	// Optionally, try to compute the command summary.
 	if output, overflow := buffer.Export(); !overflow || len(output) != 0 {
 		summaryReq := summaryRequest{
 			hosts:          hosts,
@@ -686,6 +673,13 @@ func (t *commandHandler) streamOutput(ctx context.Context, tc *client.TeleportCl
 		t.writeError(err)
 		return
 	}
+
+	if err := t.stream.Close(); err != nil {
+		t.log.WithError(err).Error("Unable to send close event to web client.")
+		return
+	}
+
+	t.log.Debug("Sent close event to web client.")
 }
 
 // connectToNodeWithMFA attempts to perform the mfa ceremony and then dial the

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

@@ -24,13 +24,19 @@ import React, {
   useRef,
 } from 'react';
 
+import type { AssistState } from 'teleport/Assist/context/state';
 import { AssistStateActionType, reducer } from 'teleport/Assist/context/state';
 
 import { convertServerMessages } from 'teleport/Assist/context/utils';
 import useStickyClusterId from 'teleport/useStickyClusterId';
 import cfg from 'teleport/config';
 import { getAccessToken, getHostName } from 'teleport/services/api';
 
+import type {
+  ConversationMessage,
+  ResolvedServerMessage,
+  ServerMessage,
+} from 'teleport/Assist/types';
 import {
   ExecutionEnvelopeType,
   RawPayload,
@@ -47,13 +53,6 @@ import {
 import * as service from '../service';
 import { resolveServerCommandMessage, resolveServerMessage } from '../service';
 
-import type {
-  ConversationMessage,
-  ResolvedServerMessage,
-  ServerMessage,
-} from 'teleport/Assist/types';
-import type { AssistState } from 'teleport/Assist/context/state';
-
 interface AssistContextValue {
   cancelMfaChallenge: () => void;
   createConversation: () => Promise<string>;
@@ -413,6 +412,8 @@ export function AssistContextProvider(props: PropsWithChildren<unknown>) {
     executeCommandWebSocket.current = new WebSocket(url);
     executeCommandWebSocket.current.binaryType = 'arraybuffer';
 
+    let sessionsEnded = 0;
+
     executeCommandWebSocket.current.onmessage = event => {
       const uintArray = new Uint8Array(event.data);
 
@@ -460,18 +461,28 @@ 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) {
+            const message = proto.encodeCloseMessage();
+            const bytearray = new Uint8Array(message);
+
+            for (const nodeId of nodeIdToResultId.keys()) {
+              dispatch({
+                type: AssistStateActionType.FinishCommandResult,
+                conversationId: state.conversations.selectedId,
+                commandResultId: nodeIdToResultId.get(nodeId),
+              });
+
+              executeCommandWebSocket.current.send(bytearray.buffer);
+            }
+
+            nodeIdToResultId.clear();
           }
 
-          nodeIdToResultId.clear();
-
-          executeCommandWebSocket.current.close();
-
           break;
       }
     };

web/packages/teleport/src/lib/term/protobuf.js

@@ -90,6 +90,12 @@ export class Protobuf {
     return this.encode(messageFields.type.values.data, message);
   }
 
+  encodeCloseMessage() {
+    // Close message has no payload
+    return this.encode(messageFields.type.values.close, '');
+  }
+
+
   encodePayload(buffer, text) {
     // set type
     buffer.push(messageFields.payload.code);