From 143cad89834e2225de2b9a94f2315be6571354d4 Mon Sep 17 00:00:00 2001
From: JerrettDavis <mxjerrett@gmail.com>
Date: Mon, 1 Jun 2026 22:05:09 -0500
Subject: [PATCH] refactor: reuse MCP discovery chain

---
 .../Registry/McpRegistry.cs                   | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/JD.SemanticKernel.Extensions.Mcp/Registry/McpRegistry.cs b/src/JD.SemanticKernel.Extensions.Mcp/Registry/McpRegistry.cs
index 92d9bea..c488de3 100644
--- a/src/JD.SemanticKernel.Extensions.Mcp/Registry/McpRegistry.cs
+++ b/src/JD.SemanticKernel.Extensions.Mcp/Registry/McpRegistry.cs
@@ -18,6 +18,7 @@ namespace JD.SemanticKernel.Extensions.Mcp.Registry;
 public sealed class McpRegistry : IMcpRegistry
 {
     private readonly IReadOnlyList<IMcpDiscoveryProvider> _providers;
+    private readonly AsyncActionChain<McpDiscoveryState> _discoveryChain;
 
     /// <summary>
     /// Initializes a new instance of <see cref="McpRegistry"/>.
@@ -31,6 +32,7 @@ public McpRegistry(IReadOnlyList<IMcpDiscoveryProvider> providers)
         if (providers is null) throw new ArgumentNullException(nameof(providers));
 #endif
         _providers = providers;
+        _discoveryChain = BuildDiscoveryChain(_providers);
     }
 
     /// <inheritdoc/>
@@ -39,11 +41,20 @@ public async Task<IReadOnlyList<McpServerDefinition>> GetAllAsync(
     {
         var merged = new Dictionary<string, McpServerDefinition>(StringComparer.OrdinalIgnoreCase);
         var state = new McpDiscoveryState(merged);
-        var chainBuilder = AsyncActionChain<McpDiscoveryState>.Create();
 
-        foreach (var provider in _providers)
+        await _discoveryChain.ExecuteAsync(state, cancellationToken).ConfigureAwait(false);
+
+        return new List<McpServerDefinition>(merged.Values).AsReadOnly();
+    }
+
+    private static AsyncActionChain<McpDiscoveryState> BuildDiscoveryChain(
+        IReadOnlyList<IMcpDiscoveryProvider> providers)
+    {
+        var builder = AsyncActionChain<McpDiscoveryState>.Create();
+
+        foreach (var provider in providers)
         {
-            chainBuilder.Use(async (current, token, next) =>
+            builder.Use(async (current, token, next) =>
             {
                 token.ThrowIfCancellationRequested();
                 var servers = await provider.DiscoverAsync(token).ConfigureAwait(false);
@@ -52,10 +63,7 @@ public async Task<IReadOnlyList<McpServerDefinition>> GetAllAsync(
             });
         }
 
-        var chain = chainBuilder.Build();
-        await chain.ExecuteAsync(state, cancellationToken).ConfigureAwait(false);
-
-        return new List<McpServerDefinition>(merged.Values).AsReadOnly();
+        return builder.Build();
     }
 
     /// <inheritdoc/>