From 81d3f62b0e7891593a0b2934025cfdeea51d5982 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 18 Jun 2025 13:02:54 +0000 Subject: [PATCH 1/3] fix(cli)!: select sub agent if available and error on multiple agents 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), it was decided prioritize them if present. Similarly as before, selecting a devcontainer randomly would be confusing. We now error out if the agent name is not specified and there are multiple agents. Fixes coder/internal#696 --- cli/ssh.go | 42 +++++++++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index 4adbf12cccf7e..5908689aa45e0 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -925,36 +925,44 @@ 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 + subAgents []codersdk.WorkspaceAgent + ) for _, resource := range resources { - agents = append(agents, resource.Agents...) + for _, agent := range resource.Agents { + availableNames = append(availableNames, agent.Name) + if agent.ParentID.UUID == uuid.Nil { + agents = append(agents, agent) + } else { + subAgents = append(subAgents, agent) + } + } } if len(agents) == 0 { return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name) } + slices.Sort(availableNames) if agentName != "" { + agents = append(agents, subAgents...) 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(subAgents) == 1 { + return subAgents[0], nil + } else if len(subAgents) > 1 { + return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple sub-agents found, please specify the agent name, available agents: %v", availableNames) + } + 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 From 716a2fba3a315ee4e46c2e1168bc54a89b68f0bb Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Wed, 18 Jun 2025 15:02:34 +0000 Subject: [PATCH 2/3] add tests --- cli/ssh_internal_test.go | 193 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 193 insertions(+) diff --git a/cli/ssh_internal_test.go b/cli/ssh_internal_test.go index 003bc697a4052..e7659dfcaa699 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,195 @@ 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, + ParentID: uuid.NullUUID{}, + } + } + + createSubAgent := func(name string, parentID uuid.UUID) codersdk.WorkspaceAgent { + return codersdk.WorkspaceAgent{ + ID: uuid.New(), + Name: name, + ParentID: uuid.NullUUID{ + UUID: parentID, + Valid: true, + }, + } + } + + 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("SingleSubAgent_NoNameSpecified", func(t *testing.T) { + t.Parallel() + parentAgent := createAgent("main") + subAgent := createSubAgent("devcontainer", parentAgent.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent}) + + // Should prefer the sub-agent when no name is specified. + result, err := getWorkspaceAgent(workspace, "") + require.NoError(t, err) + assert.Equal(t, subAgent.ID, result.ID) + assert.Equal(t, "devcontainer", result.Name) + }) + + t.Run("MultipleAgents_NoSubAgents_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("MultipleSubAgents_NoNameSpecified", func(t *testing.T) { + t.Parallel() + parentAgent := createAgent("main") + subAgent1 := createSubAgent("devcontainer1", parentAgent.ID) + subAgent2 := createSubAgent("devcontainer2", parentAgent.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent1, subAgent2}) + + _, err := getWorkspaceAgent(workspace, "") + require.Error(t, err) + assert.Contains(t, err.Error(), "multiple sub-agents found") + assert.Contains(t, err.Error(), "available agents: [devcontainer1 devcontainer2 main]") + }) + + t.Run("AgentNameSpecified_Found_RegularAgent", 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_Found_SubAgent", func(t *testing.T) { + t.Parallel() + agent := createAgent("main") + subAgent := createSubAgent("devcontainer", agent.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) + + result, err := getWorkspaceAgent(workspace, "devcontainer") + require.NoError(t, err) + assert.Equal(t, subAgent.ID, result.ID) + assert.Equal(t, "devcontainer", 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("MixedAgents_SubAgentPreferred", func(t *testing.T) { + t.Parallel() + agent := createAgent("main") + subAgent := createSubAgent("devcontainer", agent.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) + + // When no name is specified and there's one sub-agent, + // it should be preferred. + result, err := getWorkspaceAgent(workspace, "") + require.NoError(t, err) + assert.Equal(t, subAgent.ID, result.ID) + assert.Equal(t, "devcontainer", result.Name) + }) + + t.Run("MixedAgents_SpecificNameFound", func(t *testing.T) { + t.Parallel() + agent := createAgent("main") + subAgent := createSubAgent("devcontainer", agent.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) + + // Should be able to find regular agent by name. + result, err := getWorkspaceAgent(workspace, "main") + require.NoError(t, err) + assert.Equal(t, agent.ID, result.ID) + assert.Equal(t, "main", result.Name) + + // Should be able to find sub-agent by name. + result, err = getWorkspaceAgent(workspace, "devcontainer") + require.NoError(t, err) + assert.Equal(t, subAgent.ID, result.ID) + assert.Equal(t, "devcontainer", result.Name) + }) + + t.Run("AvailableAgentNames_SortedCorrectly", func(t *testing.T) { + t.Parallel() + // Define agents in non-alphabetical order. + agent2 := createAgent("zod") + agent1 := createAgent("clark") + subAgent := createSubAgent("krypton", agent1.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent2, agent1, subAgent}) + + _, err := getWorkspaceAgent(workspace, "nonexistent") + require.Error(t, err) + // Available agents should be sorted alphabetically. + assert.Contains(t, err.Error(), "available agents: [clark krypton zod]") + }) + + t.Run("MultipleAgentsAndSubAgents_NoNameSpecified", func(t *testing.T) { + t.Parallel() + agent1 := createAgent("main1") + agent2 := createAgent("main2") + subAgent1 := createSubAgent("dev1", agent1.ID) + subAgent2 := createSubAgent("dev2", agent1.ID) + workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2, subAgent1, subAgent2}) + + // Should error because there are multiple sub-agents. + _, err := getWorkspaceAgent(workspace, "") + require.Error(t, err) + assert.Contains(t, err.Error(), "multiple sub-agents found") + assert.Contains(t, err.Error(), "available agents: [dev1 dev2 main1 main2]") + }) +} From 8319ea6f85ddc9ce0ea660195feee94b413b8e5b Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Thu, 19 Jun 2025 15:35:59 +0000 Subject: [PATCH 3/3] simplify according to PR feedback --- cli/ssh.go | 13 +---- cli/ssh_internal_test.go | 106 ++------------------------------------- 2 files changed, 5 insertions(+), 114 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index 5908689aa45e0..56ab0b2a0d3af 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -928,16 +928,11 @@ func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspac var ( availableNames []string agents []codersdk.WorkspaceAgent - subAgents []codersdk.WorkspaceAgent ) for _, resource := range resources { for _, agent := range resource.Agents { availableNames = append(availableNames, agent.Name) - if agent.ParentID.UUID == uuid.Nil { - agents = append(agents, agent) - } else { - subAgents = append(subAgents, agent) - } + agents = append(agents, agent) } } if len(agents) == 0 { @@ -945,7 +940,6 @@ func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspac } slices.Sort(availableNames) if agentName != "" { - agents = append(agents, subAgents...) for _, otherAgent := range agents { if otherAgent.Name != agentName { continue @@ -954,11 +948,6 @@ func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspac } return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q, available agents: %v", agentName, availableNames) } - if len(subAgents) == 1 { - return subAgents[0], nil - } else if len(subAgents) > 1 { - return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple sub-agents found, please specify the agent name, available agents: %v", availableNames) - } if len(agents) == 1 { return agents[0], nil } diff --git a/cli/ssh_internal_test.go b/cli/ssh_internal_test.go index e7659dfcaa699..0d956def68938 100644 --- a/cli/ssh_internal_test.go +++ b/cli/ssh_internal_test.go @@ -365,21 +365,9 @@ func Test_getWorkspaceAgent(t *testing.T) { } createAgent := func(name string) codersdk.WorkspaceAgent { - return codersdk.WorkspaceAgent{ - ID: uuid.New(), - Name: name, - ParentID: uuid.NullUUID{}, - } - } - - createSubAgent := func(name string, parentID uuid.UUID) codersdk.WorkspaceAgent { return codersdk.WorkspaceAgent{ ID: uuid.New(), Name: name, - ParentID: uuid.NullUUID{ - UUID: parentID, - Valid: true, - }, } } @@ -394,20 +382,7 @@ func Test_getWorkspaceAgent(t *testing.T) { assert.Equal(t, "main", result.Name) }) - t.Run("SingleSubAgent_NoNameSpecified", func(t *testing.T) { - t.Parallel() - parentAgent := createAgent("main") - subAgent := createSubAgent("devcontainer", parentAgent.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent}) - - // Should prefer the sub-agent when no name is specified. - result, err := getWorkspaceAgent(workspace, "") - require.NoError(t, err) - assert.Equal(t, subAgent.ID, result.ID) - assert.Equal(t, "devcontainer", result.Name) - }) - - t.Run("MultipleAgents_NoSubAgents_NoNameSpecified", func(t *testing.T) { + t.Run("MultipleAgents_NoNameSpecified", func(t *testing.T) { t.Parallel() agent1 := createAgent("main1") agent2 := createAgent("main2") @@ -419,20 +394,7 @@ func Test_getWorkspaceAgent(t *testing.T) { assert.Contains(t, err.Error(), "available agents: [main1 main2]") }) - t.Run("MultipleSubAgents_NoNameSpecified", func(t *testing.T) { - t.Parallel() - parentAgent := createAgent("main") - subAgent1 := createSubAgent("devcontainer1", parentAgent.ID) - subAgent2 := createSubAgent("devcontainer2", parentAgent.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent1, subAgent2}) - - _, err := getWorkspaceAgent(workspace, "") - require.Error(t, err) - assert.Contains(t, err.Error(), "multiple sub-agents found") - assert.Contains(t, err.Error(), "available agents: [devcontainer1 devcontainer2 main]") - }) - - t.Run("AgentNameSpecified_Found_RegularAgent", func(t *testing.T) { + t.Run("AgentNameSpecified_Found", func(t *testing.T) { t.Parallel() agent1 := createAgent("main1") agent2 := createAgent("main2") @@ -444,18 +406,6 @@ func Test_getWorkspaceAgent(t *testing.T) { assert.Equal(t, "main1", result.Name) }) - t.Run("AgentNameSpecified_Found_SubAgent", func(t *testing.T) { - t.Parallel() - agent := createAgent("main") - subAgent := createSubAgent("devcontainer", agent.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) - - result, err := getWorkspaceAgent(workspace, "devcontainer") - require.NoError(t, err) - assert.Equal(t, subAgent.ID, result.ID) - assert.Equal(t, "devcontainer", result.Name) - }) - t.Run("AgentNameSpecified_NotFound", func(t *testing.T) { t.Parallel() agent1 := createAgent("main1") @@ -477,65 +427,17 @@ func Test_getWorkspaceAgent(t *testing.T) { assert.Contains(t, err.Error(), `workspace "test-workspace" has no agents`) }) - t.Run("MixedAgents_SubAgentPreferred", func(t *testing.T) { - t.Parallel() - agent := createAgent("main") - subAgent := createSubAgent("devcontainer", agent.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) - - // When no name is specified and there's one sub-agent, - // it should be preferred. - result, err := getWorkspaceAgent(workspace, "") - require.NoError(t, err) - assert.Equal(t, subAgent.ID, result.ID) - assert.Equal(t, "devcontainer", result.Name) - }) - - t.Run("MixedAgents_SpecificNameFound", func(t *testing.T) { - t.Parallel() - agent := createAgent("main") - subAgent := createSubAgent("devcontainer", agent.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent}) - - // Should be able to find regular agent by name. - result, err := getWorkspaceAgent(workspace, "main") - require.NoError(t, err) - assert.Equal(t, agent.ID, result.ID) - assert.Equal(t, "main", result.Name) - - // Should be able to find sub-agent by name. - result, err = getWorkspaceAgent(workspace, "devcontainer") - require.NoError(t, err) - assert.Equal(t, subAgent.ID, result.ID) - assert.Equal(t, "devcontainer", result.Name) - }) - t.Run("AvailableAgentNames_SortedCorrectly", func(t *testing.T) { t.Parallel() // Define agents in non-alphabetical order. agent2 := createAgent("zod") agent1 := createAgent("clark") - subAgent := createSubAgent("krypton", agent1.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent2, agent1, subAgent}) + 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]") }) - - t.Run("MultipleAgentsAndSubAgents_NoNameSpecified", func(t *testing.T) { - t.Parallel() - agent1 := createAgent("main1") - agent2 := createAgent("main2") - subAgent1 := createSubAgent("dev1", agent1.ID) - subAgent2 := createSubAgent("dev2", agent1.ID) - workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2, subAgent1, subAgent2}) - - // Should error because there are multiple sub-agents. - _, err := getWorkspaceAgent(workspace, "") - require.Error(t, err) - assert.Contains(t, err.Error(), "multiple sub-agents found") - assert.Contains(t, err.Error(), "available agents: [dev1 dev2 main1 main2]") - }) }