simplify according to PR feedback

mafredri · mafredri · commit 8319ea6f85dd · 2025-06-19T15:35:59.000Z
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -928,24 +928,18 @@ 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 {
 		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
@@ -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
@@ -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]")
-	})
 }

Original file line number	Diff line number	Diff line change
`@@ -928,24 +928,18 @@ func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspac`
`928`	`928`	`var (`
`929`	`929`	`availableNames []string`
`930`	`930`	`agents []codersdk.WorkspaceAgent`
`931`		`- subAgents []codersdk.WorkspaceAgent`
`932`	`931`	`)`
`933`	`932`	`for _, resource := range resources {`
`934`	`933`	`for _, agent := range resource.Agents {`
`935`	`934`	`availableNames = append(availableNames, agent.Name)`
`936`		`- if agent.ParentID.UUID == uuid.Nil {`
`937`		`- agents = append(agents, agent)`
`938`		`- } else {`
`939`		`- subAgents = append(subAgents, agent)`
`940`		`- }`
	`935`	`+ agents = append(agents, agent)`
`941`	`936`	`}`
`942`	`937`	`}`
`943`	`938`	`if len(agents) == 0 {`
`944`	`939`	`return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)`
`945`	`940`	`}`
`946`	`941`	`slices.Sort(availableNames)`
`947`	`942`	`if agentName != "" {`
`948`		`- agents = append(agents, subAgents...)`
`949`	`943`	`for _, otherAgent := range agents {`
`950`	`944`	`if otherAgent.Name != agentName {`
`951`	`945`	`continue`
`@@ -954,11 +948,6 @@ func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspac`
`954`	`948`	`}`
`955`	`949`	`return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q, available agents: %v", agentName, availableNames)`
`956`	`950`	`}`
`957`		`- if len(subAgents) == 1 {`
`958`		`- return subAgents[0], nil`
`959`		`- } else if len(subAgents) > 1 {`
`960`		`- return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple sub-agents found, please specify the agent name, available agents: %v", availableNames)`
`961`		`- }`
`962`	`951`	`if len(agents) == 1 {`
`963`	`952`	`return agents[0], nil`
`964`	`953`	`}`