diff --git a/lib/ai/chat_test.go b/lib/ai/chat_test.go
index e571be3279b57..1b25591450125 100644
--- a/lib/ai/chat_test.go
+++ b/lib/ai/chat_test.go
@@ -49,7 +49,7 @@ func TestChat_PromptTokens(t *testing.T) {
 					Content: "Hello",
 				},
 			},
-			want: 703,
+			want: 743,
 		},
 		{
 			name: "system and user messages",
@@ -63,7 +63,7 @@ func TestChat_PromptTokens(t *testing.T) {
 					Content: "Hi LLM.",
 				},
 			},
-			want: 711,
+			want: 751,
 		},
 		{
 			name: "tokenize our prompt",
@@ -77,7 +77,7 @@ func TestChat_PromptTokens(t *testing.T) {
 					Content: "Show me free disk space on localhost node.",
 				},
 			},
-			want: 914,
+			want: 954,
 		},
 	}
 
diff --git a/lib/ai/model/agent.go b/lib/ai/model/agent.go
index b766a47b5ddf8..37761ff940a6a 100644
--- a/lib/ai/model/agent.go
+++ b/lib/ai/model/agent.go
@@ -54,7 +54,7 @@ type Agent struct {
 	tools []Tool
 }
 
-// agentAction is an event type represetning the decision to take a single action, typically a tool invocation.
+// agentAction is an event type representing the decision to take a single action, typically a tool invocation.
 type agentAction struct {
 	// The action to take, typically a tool name.
 	action string
diff --git a/lib/ai/model/messages.go b/lib/ai/model/messages.go
index 4a4a1487adf62..26762c1bc9a84 100644
--- a/lib/ai/model/messages.go
+++ b/lib/ai/model/messages.go
@@ -67,7 +67,7 @@ type TokensUsed struct {
 }
 
 // UsedTokens returns the number of tokens used during a single invocation of the agent.
-// This method creates a convinient way to get TokensUsed from embedded structs.
+// This method creates a convenient way to get TokensUsed from embedded structs.
 func (t *TokensUsed) UsedTokens() *TokensUsed {
 	return t
 }
diff --git a/lib/ai/model/tool.go b/lib/ai/model/tool.go
index 93649f7cb6d7d..a917286eb31ab 100644
--- a/lib/ai/model/tool.go
+++ b/lib/ai/model/tool.go
@@ -53,13 +53,13 @@ func (c *commandExecutionTool) Name() string {
 }
 
 func (c *commandExecutionTool) Description() string {
-	return fmt.Sprintf(`Execute a command on a set of remote hosts based on a set of hostnames or/and a set of labels.
+	return fmt.Sprintf(`Execute a command on a set of remote nodes based on a set of node names or/and a set of labels.
 The input must be a JSON object with the following schema:
 
 %vjson
 {
 	"command": string, \\ The command to execute
-	"nodes": []string, \\ Execute a command on all nodes that have the given hostnames
+	"nodes": []string, \\ Execute a command on all nodes that have the given node names
 	"labels": []{"key": string, "value": string} \\ Execute a command on all nodes that has at least one of the labels
 }
 %v
@@ -140,7 +140,7 @@ func (e *embeddingRetrievalTool) Run(ctx context.Context, input string) (string,
 	if sb.Len() == 0 {
 		// Either no nodes are connected, embedding process hasn't started yet, or
 		// the user doesn't have access to any resources.
-		return "Didn't find any nodes matching your query", nil
+		return "Didn't find any nodes matching the query", nil
 	}
 
 	return sb.String(), nil
@@ -151,11 +151,13 @@ func (e *embeddingRetrievalTool) Name() string {
 }
 
 func (e *embeddingRetrievalTool) Description() string {
-	return fmt.Sprintf(`Ask about existing remote hosts to fetch node names or/and set of labels. Use this capability instead of guessing the names and labels.
+	return fmt.Sprintf(`Ask about existing remote nodes that user has access to fetch node names or/and set of labels. 
+Always use this capability before returning generating any command. Do not assume that the user has access to any nodes. Returning a command without checking for access will result in an error.
+Always prefer to use labler rather than node names.
 The input must be a JSON object with the following schema:
 %vjson
 {
-	"question": string \\ Question about the available remote hosts
+	"question": string \\ Question about the available remote nodes
 }
 %v
 `, "```", "```")
diff --git a/web/packages/teleport/src/Assist/context/AssistContext.tsx b/web/packages/teleport/src/Assist/context/AssistContext.tsx
index b61bb8de0535c..c18114ad74cd5 100644
--- a/web/packages/teleport/src/Assist/context/AssistContext.tsx
+++ b/web/packages/teleport/src/Assist/context/AssistContext.tsx
@@ -489,6 +489,21 @@ export function AssistContextProvider(props: PropsWithChildren<unknown>) {
           break;
       }
     };
+
+    executeCommandWebSocket.current.onclose = () => {
+      executeCommandWebSocket.current = null;
+
+      // If the execution failed, we won't get a SESSION_END message, so we
+      // need to mark all the results as finished here.
+      for (const nodeId of nodeIdToResultId.keys()) {
+        dispatch({
+          type: AssistStateActionType.FinishCommandResult,
+          conversationId: state.conversations.selectedId,
+          commandResultId: nodeIdToResultId.get(nodeId),
+        });
+      }
+      nodeIdToResultId.clear();
+    };
   }
 
   async function deleteConversation(conversationId: string) {