From bacdc28881bd35b45afcc40d5466faaf829e0a3a Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 19 Jun 2025 22:13:35 +0300 Subject: [PATCH] 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 | 31 ++++++------- cli/ssh_internal_test.go | 95 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 17 deletions(-) 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]") + }) +}