@@ -11,6 +11,7 @@ import (
1111"time"
1212
1313gliderssh"github.com/gliderlabs/ssh"
14+ "github.com/google/uuid"
1415"github.com/stretchr/testify/assert"
1516"github.com/stretchr/testify/require"
1617"golang.org/x/crypto/ssh"
@@ -346,3 +347,195 @@ func newAsyncCloser(ctx context.Context, t *testing.T) *asyncCloser {
346347started :make (chan struct {}),
347348}
348349}
350+
351+ func Test_getWorkspaceAgent (t * testing.T ) {
352+ t .Parallel ()
353+
354+ createWorkspaceWithAgents := func (agents []codersdk.WorkspaceAgent ) codersdk.Workspace {
355+ return codersdk.Workspace {
356+ Name :"test-workspace" ,
357+ LatestBuild : codersdk.WorkspaceBuild {
358+ Resources : []codersdk.WorkspaceResource {
359+ {
360+ Agents :agents ,
361+ },
362+ },
363+ },
364+ }
365+ }
366+
367+ createAgent := func (name string ) codersdk.WorkspaceAgent {
368+ return codersdk.WorkspaceAgent {
369+ ID :uuid .New (),
370+ Name :name ,
371+ ParentID : uuid.NullUUID {},
372+ }
373+ }
374+
375+ createSubAgent := func (name string ,parentID uuid.UUID ) codersdk.WorkspaceAgent {
376+ return codersdk.WorkspaceAgent {
377+ ID :uuid .New (),
378+ Name :name ,
379+ ParentID : uuid.NullUUID {
380+ UUID :parentID ,
381+ Valid :true ,
382+ },
383+ }
384+ }
385+
386+ t .Run ("SingleAgent_NoNameSpecified" ,func (t * testing.T ) {
387+ t .Parallel ()
388+ agent := createAgent ("main" )
389+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent })
390+
391+ result ,err := getWorkspaceAgent (workspace ,"" )
392+ require .NoError (t ,err )
393+ assert .Equal (t ,agent .ID ,result .ID )
394+ assert .Equal (t ,"main" ,result .Name )
395+ })
396+
397+ t .Run ("SingleSubAgent_NoNameSpecified" ,func (t * testing.T ) {
398+ t .Parallel ()
399+ parentAgent := createAgent ("main" )
400+ subAgent := createSubAgent ("devcontainer" ,parentAgent .ID )
401+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {parentAgent ,subAgent })
402+
403+ // Should prefer the sub-agent when no name is specified.
404+ result ,err := getWorkspaceAgent (workspace ,"" )
405+ require .NoError (t ,err )
406+ assert .Equal (t ,subAgent .ID ,result .ID )
407+ assert .Equal (t ,"devcontainer" ,result .Name )
408+ })
409+
410+ t .Run ("MultipleAgents_NoSubAgents_NoNameSpecified" ,func (t * testing.T ) {
411+ t .Parallel ()
412+ agent1 := createAgent ("main1" )
413+ agent2 := createAgent ("main2" )
414+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent1 ,agent2 })
415+
416+ _ ,err := getWorkspaceAgent (workspace ,"" )
417+ require .Error (t ,err )
418+ assert .Contains (t ,err .Error (),"multiple agents found" )
419+ assert .Contains (t ,err .Error (),"available agents: [main1 main2]" )
420+ })
421+
422+ t .Run ("MultipleSubAgents_NoNameSpecified" ,func (t * testing.T ) {
423+ t .Parallel ()
424+ parentAgent := createAgent ("main" )
425+ subAgent1 := createSubAgent ("devcontainer1" ,parentAgent .ID )
426+ subAgent2 := createSubAgent ("devcontainer2" ,parentAgent .ID )
427+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {parentAgent ,subAgent1 ,subAgent2 })
428+
429+ _ ,err := getWorkspaceAgent (workspace ,"" )
430+ require .Error (t ,err )
431+ assert .Contains (t ,err .Error (),"multiple sub-agents found" )
432+ assert .Contains (t ,err .Error (),"available agents: [devcontainer1 devcontainer2 main]" )
433+ })
434+
435+ t .Run ("AgentNameSpecified_Found_RegularAgent" ,func (t * testing.T ) {
436+ t .Parallel ()
437+ agent1 := createAgent ("main1" )
438+ agent2 := createAgent ("main2" )
439+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent1 ,agent2 })
440+
441+ result ,err := getWorkspaceAgent (workspace ,"main1" )
442+ require .NoError (t ,err )
443+ assert .Equal (t ,agent1 .ID ,result .ID )
444+ assert .Equal (t ,"main1" ,result .Name )
445+ })
446+
447+ t .Run ("AgentNameSpecified_Found_SubAgent" ,func (t * testing.T ) {
448+ t .Parallel ()
449+ agent := createAgent ("main" )
450+ subAgent := createSubAgent ("devcontainer" ,agent .ID )
451+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent ,subAgent })
452+
453+ result ,err := getWorkspaceAgent (workspace ,"devcontainer" )
454+ require .NoError (t ,err )
455+ assert .Equal (t ,subAgent .ID ,result .ID )
456+ assert .Equal (t ,"devcontainer" ,result .Name )
457+ })
458+
459+ t .Run ("AgentNameSpecified_NotFound" ,func (t * testing.T ) {
460+ t .Parallel ()
461+ agent1 := createAgent ("main1" )
462+ agent2 := createAgent ("main2" )
463+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent1 ,agent2 })
464+
465+ _ ,err := getWorkspaceAgent (workspace ,"nonexistent" )
466+ require .Error (t ,err )
467+ assert .Contains (t ,err .Error (),`agent not found by name "nonexistent"` )
468+ assert .Contains (t ,err .Error (),"available agents: [main1 main2]" )
469+ })
470+
471+ t .Run ("NoAgents" ,func (t * testing.T ) {
472+ t .Parallel ()
473+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {})
474+
475+ _ ,err := getWorkspaceAgent (workspace ,"" )
476+ require .Error (t ,err )
477+ assert .Contains (t ,err .Error (),`workspace "test-workspace" has no agents` )
478+ })
479+
480+ t .Run ("MixedAgents_SubAgentPreferred" ,func (t * testing.T ) {
481+ t .Parallel ()
482+ agent := createAgent ("main" )
483+ subAgent := createSubAgent ("devcontainer" ,agent .ID )
484+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent ,subAgent })
485+
486+ // When no name is specified and there's one sub-agent,
487+ // it should be preferred.
488+ result ,err := getWorkspaceAgent (workspace ,"" )
489+ require .NoError (t ,err )
490+ assert .Equal (t ,subAgent .ID ,result .ID )
491+ assert .Equal (t ,"devcontainer" ,result .Name )
492+ })
493+
494+ t .Run ("MixedAgents_SpecificNameFound" ,func (t * testing.T ) {
495+ t .Parallel ()
496+ agent := createAgent ("main" )
497+ subAgent := createSubAgent ("devcontainer" ,agent .ID )
498+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent ,subAgent })
499+
500+ // Should be able to find regular agent by name.
501+ result ,err := getWorkspaceAgent (workspace ,"main" )
502+ require .NoError (t ,err )
503+ assert .Equal (t ,agent .ID ,result .ID )
504+ assert .Equal (t ,"main" ,result .Name )
505+
506+ // Should be able to find sub-agent by name.
507+ result ,err = getWorkspaceAgent (workspace ,"devcontainer" )
508+ require .NoError (t ,err )
509+ assert .Equal (t ,subAgent .ID ,result .ID )
510+ assert .Equal (t ,"devcontainer" ,result .Name )
511+ })
512+
513+ t .Run ("AvailableAgentNames_SortedCorrectly" ,func (t * testing.T ) {
514+ t .Parallel ()
515+ // Define agents in non-alphabetical order.
516+ agent2 := createAgent ("zod" )
517+ agent1 := createAgent ("clark" )
518+ subAgent := createSubAgent ("krypton" ,agent1 .ID )
519+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent2 ,agent1 ,subAgent })
520+
521+ _ ,err := getWorkspaceAgent (workspace ,"nonexistent" )
522+ require .Error (t ,err )
523+ // Available agents should be sorted alphabetically.
524+ assert .Contains (t ,err .Error (),"available agents: [clark krypton zod]" )
525+ })
526+
527+ t .Run ("MultipleAgentsAndSubAgents_NoNameSpecified" ,func (t * testing.T ) {
528+ t .Parallel ()
529+ agent1 := createAgent ("main1" )
530+ agent2 := createAgent ("main2" )
531+ subAgent1 := createSubAgent ("dev1" ,agent1 .ID )
532+ subAgent2 := createSubAgent ("dev2" ,agent1 .ID )
533+ workspace := createWorkspaceWithAgents ([]codersdk.WorkspaceAgent {agent1 ,agent2 ,subAgent1 ,subAgent2 })
534+
535+ // Should error because there are multiple sub-agents.
536+ _ ,err := getWorkspaceAgent (workspace ,"" )
537+ require .Error (t ,err )
538+ assert .Contains (t ,err .Error (),"multiple sub-agents found" )
539+ assert .Contains (t ,err .Error (),"available agents: [dev1 dev2 main1 main2]" )
540+ })
541+ }