fix(cli)!: enforce selection for multiple agents rather than use randomness (#18427)

In the past we randomly selected workspace agent if there were multiple.
Unless both are running on the same machine with the same configuration,
this would be very confusing behavior for a user.

With the introduction of sub agents (devcontainer agents), we have now
made this an error state and require the specifying of agent when there
is more than one (either normal agent or sub agent).

This aligns with the behavior of e.g. Coder Desktop.

Fixes coder/internal#696

cli/ssh.go

@@ -925,36 +925,33 @@ func getWorkspaceAndAgent(ctx context.Context, inv *serpent.Invocation, client *
 func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspaceAgent codersdk.WorkspaceAgent, err error) {
 	resources := workspace.LatestBuild.Resources
 
-	agents := make([]codersdk.WorkspaceAgent, 0)
+	var (
+		availableNames []string
+		agents         []codersdk.WorkspaceAgent
+	)
 	for _, resource := range resources {
-		agents = append(agents, resource.Agents...)
+		for _, agent := range resource.Agents {
+			availableNames = append(availableNames, agent.Name)
+			agents = append(agents, agent)
+		}
 	}
 	if len(agents) == 0 {
 		return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)
 	}
+	slices.Sort(availableNames)
 	if agentName != "" {
 		for _, otherAgent := range agents {
 			if otherAgent.Name != agentName {
 				continue
 			}
-			workspaceAgent = otherAgent
-			break
-		}
-		if workspaceAgent.ID == uuid.Nil {
-			return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q", agentName)
+			return otherAgent, nil
 		}
+		return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q, available agents: %v", agentName, availableNames)
 	}
-	if workspaceAgent.ID == uuid.Nil {
-		if len(agents) > 1 {
-			workspaceAgent, err = cryptorand.Element(agents)
-			if err != nil {
-				return codersdk.WorkspaceAgent{}, err
-			}
-		} else {
-			workspaceAgent = agents[0]
-		}
+	if len(agents) == 1 {
+		return agents[0], nil
 	}
-	return workspaceAgent, nil
+	return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple agents found, please specify the agent name, available agents: %v", availableNames)
 }
 
 // Attempt to poll workspace autostop. We write a per-workspace lockfile to

cli/ssh_internal_test.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	gliderssh "github.com/gliderlabs/ssh"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -346,3 +347,97 @@ func newAsyncCloser(ctx context.Context, t *testing.T) *asyncCloser {
 		started:    make(chan struct{}),
 	}
 }
+
+func Test_getWorkspaceAgent(t *testing.T) {
+	t.Parallel()
+
+	createWorkspaceWithAgents := func(agents []codersdk.WorkspaceAgent) codersdk.Workspace {
+		return codersdk.Workspace{
+			Name: "test-workspace",
+			LatestBuild: codersdk.WorkspaceBuild{
+				Resources: []codersdk.WorkspaceResource{
+					{
+						Agents: agents,
+					},
+				},
+			},
+		}
+	}
+
+	createAgent := func(name string) codersdk.WorkspaceAgent {
+		return codersdk.WorkspaceAgent{
+			ID:   uuid.New(),
+			Name: name,
+		}
+	}
+
+	t.Run("SingleAgent_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		agent := createAgent("main")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent})
+
+		result, err := getWorkspaceAgent(workspace, "")
+		require.NoError(t, err)
+		assert.Equal(t, agent.ID, result.ID)
+		assert.Equal(t, "main", result.Name)
+	})
+
+	t.Run("MultipleAgents_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "multiple agents found")
+		assert.Contains(t, err.Error(), "available agents: [main1 main2]")
+	})
+
+	t.Run("AgentNameSpecified_Found", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		result, err := getWorkspaceAgent(workspace, "main1")
+		require.NoError(t, err)
+		assert.Equal(t, agent1.ID, result.ID)
+		assert.Equal(t, "main1", result.Name)
+	})
+
+	t.Run("AgentNameSpecified_NotFound", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		_, err := getWorkspaceAgent(workspace, "nonexistent")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), `agent not found by name "nonexistent"`)
+		assert.Contains(t, err.Error(), "available agents: [main1 main2]")
+	})
+
+	t.Run("NoAgents", func(t *testing.T) {
+		t.Parallel()
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{})
+
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), `workspace "test-workspace" has no agents`)
+	})
+
+	t.Run("AvailableAgentNames_SortedCorrectly", func(t *testing.T) {
+		t.Parallel()
+		// Define agents in non-alphabetical order.
+		agent2 := createAgent("zod")
+		agent1 := createAgent("clark")
+		agent3 := createAgent("krypton")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent2, agent1, agent3})
+
+		_, err := getWorkspaceAgent(workspace, "nonexistent")
+		require.Error(t, err)
+		// Available agents should be sorted alphabetically.
+		assert.Contains(t, err.Error(), "available agents: [clark krypton zod]")
+	})
+}