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

mafredri · web-flow · commit bacdc28881bd · 2025-06-19T22:13:35.000+03:00
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
diff --git a/cli/ssh.go b/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
diff --git a/cli/ssh_internal_test.go b/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]")
+	})
+}
