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]") + }) +}
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies: