Add node name to the Assist execution result (#27635)

* Add node name to the Assist execution result

Currently, only node ID is returned on the command execution result in Assist. For better UX we want to display Node name which id more human friendly rather than a node ID which is a UUID. Adding the value to returned payload sounds cheaper than calling an API to get node names.

* Add test

* Extract commandExecResult struct

* Fix test after rebase

lib/web/command.go

@@ -67,6 +67,18 @@ type CommandRequest struct {
 	ExecutionID string `json:"execution_id"`
 }
 
+// commandExecResult is a result of a command execution.
+type commandExecResult struct {
+	// NodeID is the ID of the node where the command was executed.
+	NodeID string `json:"node_id"`
+	// NodeName is the name of the node where the command was executed.
+	NodeName string `json:"node_name"`
+	// ExecutionID is a unique ID used to identify the command execution.
+	ExecutionID string `json:"execution_id"`
+	// SessionID is the ID of the session where the command was executed.
+	SessionID string `json:"session_id"`
+}
+
 // Check checks if the request is valid.
 func (c *CommandRequest) Check() error {
 	if c.Command == "" {
@@ -226,12 +238,9 @@ func (h *Handler) executeCommand(
 		h.log.Infof("Executing command: %#v.", req)
 		httplib.MakeTracingHandler(handler, teleport.ComponentProxy).ServeHTTP(w, r)
 
-		msgPayload, err := json.Marshal(struct {
-			NodeID      string `json:"node_id"`
-			ExecutionID string `json:"execution_id"`
-			SessionID   string `json:"session_id"`
-		}{
+		msgPayload, err := json.Marshal(&commandExecResult{
 			NodeID:      host.id,
+			NodeName:    host.hostName,
 			ExecutionID: req.ExecutionID,
 			SessionID:   string(sessionData.ID),
 		})

lib/web/command_test.go

@@ -37,7 +37,11 @@ import (
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/gravitational/teleport/api/gen/proto/go/assist/v1"
+	assistlib "github.com/gravitational/teleport/lib/assist"
+	"github.com/gravitational/teleport/lib/auth"
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/session"
 	"github.com/gravitational/teleport/lib/utils"
@@ -47,7 +51,7 @@ func TestExecuteCommand(t *testing.T) {
 	t.Parallel()
 	s := newWebSuite(t)
 
-	ws, _, err := s.makeCommand(t, s.authPack(t, "foo"))
+	ws, _, err := s.makeCommand(t, s.authPack(t, "foo"), uuid.New())
 	require.NoError(t, err)
 	t.Cleanup(func() { require.NoError(t, ws.Close()) })
 
@@ -56,11 +60,74 @@ func TestExecuteCommand(t *testing.T) {
 	require.NoError(t, waitForCommandOutput(stream, "teleport"))
 }
 
-func (s *WebSuite) makeCommand(t *testing.T, pack *authPack) (*websocket.Conn, *session.Session, error) {
+func TestExecuteCommandHistory(t *testing.T) {
+	t.Parallel()
+
+	// Given
+	s := newWebSuite(t)
+	authPack := s.authPack(t, "foo")
+
+	ctx := context.Background()
+	clt, err := s.server.NewClient(auth.TestUser("foo"))
+	require.NoError(t, err)
+
+	// Create conversation, otherwise the command execution will not be saved
+	conversation, err := clt.CreateAssistantConversation(context.Background(), &assist.CreateAssistantConversationRequest{
+		Username:    "foo",
+		CreatedTime: timestamppb.Now(),
+	})
+	require.NoError(t, err)
+
+	require.NotEmpty(t, conversation.GetId())
+
+	conversationID, err := uuid.Parse(conversation.GetId())
+	require.NoError(t, err)
+
+	ws, _, err := s.makeCommand(t, authPack, conversationID)
+	require.NoError(t, err)
+
+	stream := NewWStream(ctx, ws, utils.NewLoggerForTests(), nil)
+
+	// When command executes
+	require.NoError(t, waitForCommandOutput(stream, "teleport"))
+
+	// Explecitly close the stream
+	err = stream.Close()
+	require.NoError(t, err)
+
+	// Then command execution history is saved
+	var messages *assist.GetAssistantMessagesResponse
+	// Command execution history is saved in asynchronusly, so we need to wait for it.
+	require.Eventually(t, func() bool {
+		messages, err = clt.GetAssistantMessages(ctx, &assist.GetAssistantMessagesRequest{
+			ConversationId: conversationID.String(),
+			Username:       "foo",
+		})
+		require.NoError(t, err)
+
+		return len(messages.GetMessages()) == 1
+	}, 5*time.Second, 100*time.Millisecond)
+
+	// Assert the returned message
+	msg := messages.GetMessages()[0]
+	require.Equal(t, string(assistlib.MessageKindCommandResult), msg.Type)
+	require.NotZero(t, msg.CreatedTime)
+
+	var result commandExecResult
+	err = json.Unmarshal([]byte(msg.GetPayload()), &result)
+	require.NoError(t, err)
+
+	require.NotEmpty(t, result.ExecutionID)
+	require.NotEmpty(t, result.SessionID)
+	require.Equal(t, "node", result.NodeName)
+	require.Equal(t, "node", result.NodeID)
+}
+
+func (s *WebSuite) makeCommand(t *testing.T, pack *authPack, conversationID uuid.UUID) (*websocket.Conn, *session.Session, error) {
 	req := CommandRequest{
 		Query:          fmt.Sprintf("name == \"%s\"", s.srvID),
 		Login:          pack.login,
-		ConversationID: uuid.New().String(),
+		ConversationID: conversationID.String(),
 		ExecutionID:    uuid.New().String(),
 		Command:        "echo txlxport | sed 's/x/e/g'",
 	}