fix: resolve sub-workflow agent discovery

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

sidecar/src/Aryx.AgentHost/Services/CopilotAgentBundle.cs

@@ -57,7 +57,7 @@ internal sealed class CopilotAgentBundle : IAsyncDisposable
         CopilotClient sharedClient = new(clientOptions);
         await sharedClient.StartAsync(cancellationToken).ConfigureAwait(false);
 
-        IReadOnlyList<WorkflowNodeDto> agentNodes = command.Workflow.GetAgentNodes();
+        IReadOnlyList<WorkflowNodeDto> agentNodes = command.Workflow.GetAllAgentNodes(command.WorkflowLibrary);
         foreach ((WorkflowNodeDto definition, int agentIndex) in agentNodes.Select((definition, index) => (definition, index)))
         {
             SessionConfig sessionConfig = CreateSessionConfig(

sidecar/src/Aryx.AgentHost/Services/WorkflowDefinitionExtensions.cs

@@ -13,6 +13,32 @@ internal static class WorkflowDefinitionExtensions
             .ToList();
     }
 
+    public static IReadOnlyList<WorkflowNodeDto> GetAllAgentNodes(
+        this WorkflowDefinitionDto workflow,
+        IReadOnlyList<WorkflowDefinitionDto>? workflowLibrary = null)
+    {
+        ArgumentNullException.ThrowIfNull(workflow);
+
+        Dictionary<string, WorkflowDefinitionDto> workflowLibraryMap = CreateWorkflowLibraryMap(workflowLibrary);
+        return GetAllAgentNodes(workflow, workflowLibraryMap);
+    }
+
+    public static IReadOnlyList<WorkflowNodeDto> GetAllAgentNodes(
+        this WorkflowDefinitionDto workflow,
+        IReadOnlyDictionary<string, WorkflowDefinitionDto>? workflowLibrary)
+    {
+        ArgumentNullException.ThrowIfNull(workflow);
+
+        List<WorkflowNodeDto> agentNodes = [];
+        CollectAgentNodes(
+            workflow,
+            workflowLibrary ?? EmptyWorkflowLibrary,
+            agentNodes,
+            new HashSet<string>(StringComparer.Ordinal),
+            new HashSet<WorkflowDefinitionDto>(ReferenceEqualityComparer.Instance));
+        return agentNodes;
+    }
+
     public static bool IsAgentNode(this WorkflowNodeDto node)
     {
         ArgumentNullException.ThrowIfNull(node);
@@ -31,6 +57,28 @@ internal static class WorkflowDefinitionExtensions
         return FirstNonBlank(node.Config.Name, node.Label, node.Id) ?? "agent";
     }
 
+    public static WorkflowDefinitionDto ResolveSubWorkflowDefinition(
+        this WorkflowNodeDto node,
+        IReadOnlyDictionary<string, WorkflowDefinitionDto>? workflowLibrary)
+    {
+        ArgumentNullException.ThrowIfNull(node);
+
+        if (node.Config.InlineWorkflow is not null)
+        {
+            return node.Config.InlineWorkflow;
+        }
+
+        if (!string.IsNullOrWhiteSpace(node.Config.WorkflowId)
+            && workflowLibrary is not null
+            && workflowLibrary.TryGetValue(node.Config.WorkflowId, out WorkflowDefinitionDto? workflow))
+        {
+            return workflow;
+        }
+
+        throw new InvalidOperationException(
+            $"Sub-workflow node \"{node.Id}\" references unknown workflow \"{node.Config.WorkflowId}\".");
+    }
+
     public static bool IsOrchestrationMode(this WorkflowDefinitionDto workflow, string mode)
     {
         ArgumentNullException.ThrowIfNull(workflow);
@@ -39,6 +87,60 @@ internal static class WorkflowDefinitionExtensions
         return string.Equals(workflow.Settings.OrchestrationMode, mode, StringComparison.OrdinalIgnoreCase);
     }
 
+    private static readonly IReadOnlyDictionary<string, WorkflowDefinitionDto> EmptyWorkflowLibrary =
+        new Dictionary<string, WorkflowDefinitionDto>(StringComparer.Ordinal);
+
+    private static void CollectAgentNodes(
+        WorkflowDefinitionDto workflowDefinition,
+        IReadOnlyDictionary<string, WorkflowDefinitionDto> workflowLibrary,
+        List<WorkflowNodeDto> agentNodes,
+        ISet<string> visitedWorkflowIds,
+        ISet<WorkflowDefinitionDto> visitedAnonymousWorkflows)
+    {
+        string? workflowId = NormalizeOptionalString(workflowDefinition.Id);
+        if (workflowId is not null)
+        {
+            if (!visitedWorkflowIds.Add(workflowId))
+            {
+                return;
+            }
+        }
+        else if (!visitedAnonymousWorkflows.Add(workflowDefinition))
+        {
+            return;
+        }
+
+        foreach (WorkflowNodeDto node in workflowDefinition.Graph.Nodes)
+        {
+            if (node.IsAgentNode())
+            {
+                agentNodes.Add(node);
+                continue;
+            }
+
+            if (!string.Equals(node.Kind, "sub-workflow", StringComparison.OrdinalIgnoreCase))
+            {
+                continue;
+            }
+
+            WorkflowDefinitionDto subWorkflow = node.ResolveSubWorkflowDefinition(workflowLibrary);
+            CollectAgentNodes(subWorkflow, workflowLibrary, agentNodes, visitedWorkflowIds, visitedAnonymousWorkflows);
+        }
+    }
+
+    private static Dictionary<string, WorkflowDefinitionDto> CreateWorkflowLibraryMap(
+        IReadOnlyList<WorkflowDefinitionDto>? workflowLibrary)
+    {
+        return workflowLibrary?
+            .Where(candidate => !string.IsNullOrWhiteSpace(candidate.Id))
+            .GroupBy(candidate => candidate.Id, StringComparer.Ordinal)
+            .ToDictionary(group => group.Key, group => group.Last(), StringComparer.Ordinal)
+            ?? new Dictionary<string, WorkflowDefinitionDto>(StringComparer.Ordinal);
+    }
+
+    private static string? NormalizeOptionalString(string? value)
+        => string.IsNullOrWhiteSpace(value) ? null : value.Trim();
+
     private static string? FirstNonBlank(params string?[] values)
     {
         foreach (string? value in values)

sidecar/src/Aryx.AgentHost/Services/WorkflowRunner.cs

@@ -171,7 +171,7 @@ internal sealed class WorkflowRunner
 
         if (string.Equals(node.Kind, "sub-workflow", StringComparison.OrdinalIgnoreCase))
         {
-            WorkflowDefinitionDto subWorkflowDefinition = ResolveSubWorkflowDefinition(node, workflowLibrary);
+            WorkflowDefinitionDto subWorkflowDefinition = node.ResolveSubWorkflowDefinition(workflowLibrary);
             Workflow subWorkflow = BuildWorkflow(subWorkflowDefinition, agentMap, workflowLibrary);
             return new WorkflowNodeRoute(subWorkflow.BindAsExecutor(node.Id));
         }
@@ -196,25 +196,6 @@ internal sealed class WorkflowRunner
         return new WorkflowNodeRoute(entry, exit, [(entry, portBinding), (portBinding, exit)]);
     }
 
-    private static WorkflowDefinitionDto ResolveSubWorkflowDefinition(
-        WorkflowNodeDto node,
-        IReadOnlyDictionary<string, WorkflowDefinitionDto> workflowLibrary)
-    {
-        if (node.Config.InlineWorkflow is not null)
-        {
-            return node.Config.InlineWorkflow;
-        }
-
-        if (!string.IsNullOrWhiteSpace(node.Config.WorkflowId)
-            && workflowLibrary.TryGetValue(node.Config.WorkflowId, out WorkflowDefinitionDto? workflow))
-        {
-            return workflow;
-        }
-
-        throw new InvalidOperationException(
-            $"Sub-workflow node \"{node.Id}\" references unknown workflow \"{node.Config.WorkflowId}\".");
-    }
-
     private static string NormalizeRequired(string? value, string errorMessage)
         => NormalizeOptionalString(value) ?? throw new InvalidOperationException(errorMessage);
 
@@ -225,41 +206,9 @@ internal sealed class WorkflowRunner
         WorkflowDefinitionDto workflowDefinition,
         IReadOnlyDictionary<string, WorkflowDefinitionDto> workflowLibrary)
     {
-        List<string> agentIds = [];
-        CollectAgentIds(workflowDefinition, workflowLibrary, agentIds, new HashSet<string>(StringComparer.Ordinal));
-        return agentIds;
-    }
-
-    private static void CollectAgentIds(
-        WorkflowDefinitionDto workflowDefinition,
-        IReadOnlyDictionary<string, WorkflowDefinitionDto> workflowLibrary,
-        List<string> agentIds,
-        ISet<string> visitedWorkflowIds)
-    {
-        string workflowKey = string.IsNullOrWhiteSpace(workflowDefinition.Id)
-            ? Guid.NewGuid().ToString("N")
-            : workflowDefinition.Id;
-        if (!visitedWorkflowIds.Add(workflowKey))
-        {
-            return;
-        }
-
-        foreach (WorkflowNodeDto node in workflowDefinition.Graph.Nodes)
-        {
-            if (string.Equals(node.Kind, "agent", StringComparison.OrdinalIgnoreCase))
-            {
-                agentIds.Add(!string.IsNullOrWhiteSpace(node.Config.Id) ? node.Config.Id : node.Id);
-                continue;
-            }
-
-            if (!string.Equals(node.Kind, "sub-workflow", StringComparison.OrdinalIgnoreCase))
-            {
-                continue;
-            }
-
-            WorkflowDefinitionDto subWorkflow = ResolveSubWorkflowDefinition(node, workflowLibrary);
-            CollectAgentIds(subWorkflow, workflowLibrary, agentIds, visitedWorkflowIds);
-        }
+        return workflowDefinition.GetAllAgentNodes(workflowLibrary)
+            .Select(node => node.GetAgentId())
+            .ToList();
     }
 
     private sealed record WorkflowNodeRoute(