diff --git a/cli/ssh.go b/cli/ssh.go index 4adbf12cccf7e..56ab0b2a0d3af 100644 --- 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 index 003bc697a4052..0d956def68938 100644 --- 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]") + }) +}