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