From 44991a660ee42d15cd943427447abd46e11f2b5b Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@goteleport.com>
Date: Fri, 26 May 2023 16:33:26 +0300
Subject: [PATCH 1/5] Fix display of Assist executions with empty output

---
 .../Assist/Chat/ChatItem/Action/RunAction.tsx | 38 ++++++++++++++++---
 .../src/Assist/Chat/ChatItem/ChatItem.tsx     |  2 +
 .../teleport/src/Assist/contexts/messages.tsx | 23 +++++++----
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
index f12a2fbaadd1d..e09595034ec7e 100644
--- a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
+++ b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
@@ -75,6 +75,10 @@ interface RawPayload {
   payload: string;
 }
 
+interface SessionData {
+  session: { server_id: string };
+}
+
 class assistClient extends EventEmitterWebAuthnSender {
   private readonly ws: WebSocket;
   readonly proto: Protobuf = new Protobuf();
@@ -89,11 +93,32 @@ class assistClient extends EventEmitterWebAuthnSender {
     this.ws = new WebSocket(url);
     this.ws.binaryType = 'arraybuffer';
 
+    this.ws.onclose = event => {
+      setState(state => {
+        return state.map(n => ({
+          ...n,
+          stdout: n.stdout || '',
+          status: RunActionStatus.Finished,
+        }));
+      });
+    };
+
     this.ws.onmessage = event => {
       const uintArray = new Uint8Array(event.data);
       const msg = this.proto.decode(uintArray);
 
       switch (msg.type) {
+        case MessageTypeEnum.SESSION_DATA:
+          const sessionData = JSON.parse(msg.payload) as SessionData;
+          setState(state => {
+            state.push({
+              nodeId: sessionData.session.server_id,
+              status: RunActionStatus.Connecting,
+            });
+            return state;
+          });
+          break;
+
         case MessageTypeEnum.RAW:
           const data = JSON.parse(msg.payload) as RawPayload;
           const payload = atob(data.payload);
@@ -245,11 +270,14 @@ function NodeOutput(props: NodeOutputProps) {
         </LoadingContainer>
       )}
 
-      {props.state.stdout && (
-        <NodeContent>
-          <pre>{props.state.stdout}</pre>
-        </NodeContent>
-      )}
+      {props.state.stdout !== undefined &&
+        (props.state.stdout == '' ? (
+          'Empty output.'
+        ) : (
+          <NodeContent>
+            <pre>{props.state.stdout}</pre>
+          </NodeContent>
+        ))}
     </NodeContainer>
   );
 }
diff --git a/web/packages/teleport/src/Assist/Chat/ChatItem/ChatItem.tsx b/web/packages/teleport/src/Assist/Chat/ChatItem/ChatItem.tsx
index 5aab8eeada780..f58f7daf8fcb7 100644
--- a/web/packages/teleport/src/Assist/Chat/ChatItem/ChatItem.tsx
+++ b/web/packages/teleport/src/Assist/Chat/ChatItem/ChatItem.tsx
@@ -251,6 +251,8 @@ export function ChatItem(props: ChatItemProps) {
               </MachineName>
               {props.message.content.errorMsg ? (
                 <ErrorMessage>{props.message.content.errorMsg}</ErrorMessage>
+              ) : props.message.content.payload === '' ? (
+                <p>Empty output.</p>
               ) : (
                 <Output>{props.message.content.payload}</Output>
               )}
diff --git a/web/packages/teleport/src/Assist/contexts/messages.tsx b/web/packages/teleport/src/Assist/contexts/messages.tsx
index 4b1b05cc1fb59..ca1b95b053034 100644
--- a/web/packages/teleport/src/Assist/contexts/messages.tsx
+++ b/web/packages/teleport/src/Assist/contexts/messages.tsx
@@ -221,16 +221,25 @@ async function convertServerMessage(
       sid: payload.session_id,
     });
 
-    // The offset here is set base on A/B test that was run between me, myself and I.
-    const resp = await api.fetch(sessionUrl + '/stream?offset=0&bytes=4096', {
-      Accept: 'text/plain',
-      'Content-Type': 'text/plain; charset=utf-8',
-    });
+    const sessionResp = await api.fetch(sessionUrl + '/events');
+    const sessionExists = sessionResp.status === 200;
 
     let msg;
     let errorMsg;
-    if (resp.status === 200) {
-      msg = await resp.text();
+    if (sessionExists) {
+      // The offset here is set base on A/B test that was run between me, myself and I.
+      const stream = await api.fetch(
+        sessionUrl + '/stream?offset=0&bytes=4096',
+        {
+          Accept: 'text/plain',
+          'Content-Type': 'text/plain; charset=utf-8',
+        }
+      );
+      if (stream.status === 200) {
+        msg = await stream.text();
+      } else {
+        msg = ''; // Empty output, handled in <Output>
+      }
     } else {
       errorMsg = 'No session recording. The command execution failed.';
     }

From 13ec82b658e057e92287bf0136b0b79f4382ce0d Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@goteleport.com>
Date: Mon, 29 May 2023 14:59:34 +0300
Subject: [PATCH 2/5] Lint

---
 .../teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
index e09595034ec7e..6f86d014d1247 100644
--- a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
+++ b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
@@ -93,7 +93,7 @@ class assistClient extends EventEmitterWebAuthnSender {
     this.ws = new WebSocket(url);
     this.ws.binaryType = 'arraybuffer';
 
-    this.ws.onclose = event => {
+    this.ws.onclose = () => {
       setState(state => {
         return state.map(n => ({
           ...n,

From b13c6f5400bde1233071a67fe524eb0aab26af8e Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@goteleport.com>
Date: Mon, 29 May 2023 15:01:57 +0300
Subject: [PATCH 3/5] Nit

---
 .../teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
index 6f86d014d1247..d65a9babd04da 100644
--- a/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
+++ b/web/packages/teleport/src/Assist/Chat/ChatItem/Action/RunAction.tsx
@@ -271,7 +271,7 @@ function NodeOutput(props: NodeOutputProps) {
       )}
 
       {props.state.stdout !== undefined &&
-        (props.state.stdout == '' ? (
+        (props.state.stdout === '' ? (
           'Empty output.'
         ) : (
           <NodeContent>

From db5faf04a4e8375b059272ac66036223601c2646 Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@goteleport.com>
Date: Mon, 29 May 2023 15:27:11 +0300
Subject: [PATCH 4/5] Improve display of commands that failed with code

---
 .../teleport/src/Assist/contexts/messages.tsx | 22 ++++++++++++++++---
 web/packages/teleport/src/lib/term/enums.ts   |  1 +
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/web/packages/teleport/src/Assist/contexts/messages.tsx b/web/packages/teleport/src/Assist/contexts/messages.tsx
index ca1b95b053034..f47cdc97934ee 100644
--- a/web/packages/teleport/src/Assist/contexts/messages.tsx
+++ b/web/packages/teleport/src/Assist/contexts/messages.tsx
@@ -37,6 +37,8 @@ import cfg from 'teleport/config';
 
 import { ApiError } from 'teleport/services/api/parseError';
 
+import { EventType } from 'teleport/lib/term/enums';
+
 import {
   Author,
   ExecuteRemoteCommandContent,
@@ -84,6 +86,12 @@ interface PartialMessagePayload {
   idx: number;
 }
 
+interface ExecEvent {
+  event: EventType.EXEC;
+  exitError?: string;
+}
+type SessionEvent = ExecEvent | { event: string };
+
 const convertToQuery = (cmd: ExecuteRemoteCommandPayload): string => {
   let query = '';
 
@@ -221,8 +229,16 @@ async function convertServerMessage(
       sid: payload.session_id,
     });
 
-    const sessionResp = await api.fetch(sessionUrl + '/events');
-    const sessionExists = sessionResp.status === 200;
+    const eventsResp = await api.fetch(sessionUrl + '/events');
+    const sessionExists = eventsResp.status === 200;
+    const eventsData = (await eventsResp.json()) as {
+      events: SessionEvent[];
+    };
+
+    function isExecEvent(e: SessionEvent): e is ExecEvent {
+      return e.event == EventType.EXEC;
+    }
+    const execEvent = eventsData.events.find(isExecEvent);
 
     let msg;
     let errorMsg;
@@ -238,7 +254,7 @@ async function convertServerMessage(
       if (stream.status === 200) {
         msg = await stream.text();
       } else {
-        msg = ''; // Empty output, handled in <Output>
+        msg = execEvent?.exitError || ''; // empty output handled in <Output>
       }
     } else {
       errorMsg = 'No session recording. The command execution failed.';
diff --git a/web/packages/teleport/src/lib/term/enums.ts b/web/packages/teleport/src/lib/term/enums.ts
index b1cd690816765..dd422f995627f 100644
--- a/web/packages/teleport/src/lib/term/enums.ts
+++ b/web/packages/teleport/src/lib/term/enums.ts
@@ -18,6 +18,7 @@ export enum EventType {
   START = 'session.start',
   JOIN = 'session.join',
   END = 'session.end',
+  EXEC = 'exec',
   PRINT = 'print',
   RESIZE = 'resize',
   FILE_TRANSFER_REQUEST = 'file_transfer_request',

From b8e5268a6cf465ab2eb173297e6f2516a6534ae1 Mon Sep 17 00:00:00 2001
From: Justinas Stankevicius <justinas@goteleport.com>
Date: Mon, 29 May 2023 19:05:33 +0300
Subject: [PATCH 5/5] Address lint

---
 web/packages/teleport/src/Assist/contexts/messages.tsx | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/web/packages/teleport/src/Assist/contexts/messages.tsx b/web/packages/teleport/src/Assist/contexts/messages.tsx
index f47cdc97934ee..064a721f50be2 100644
--- a/web/packages/teleport/src/Assist/contexts/messages.tsx
+++ b/web/packages/teleport/src/Assist/contexts/messages.tsx
@@ -169,6 +169,10 @@ export const remoteCommandToMessage = async (
   }
 };
 
+function isExecEvent(e: SessionEvent): e is ExecEvent {
+  return e.event == EventType.EXEC;
+}
+
 async function convertServerMessage(
   message: ServerMessage,
   clusterId: string
@@ -234,10 +238,6 @@ async function convertServerMessage(
     const eventsData = (await eventsResp.json()) as {
       events: SessionEvent[];
     };
-
-    function isExecEvent(e: SessionEvent): e is ExecEvent {
-      return e.event == EventType.EXEC;
-    }
     const execEvent = eventsData.events.find(isExecEvent);
 
     let msg;