fix(mcp): bound + parallelize provider discovery so one bad provider can't stall the aggregate

providers/mcp/aggregate/provider_proxy.go

@@ -44,6 +44,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -57,12 +58,26 @@ import (
 // aggregator stays decoupled from pkg/hub/providers.
 type ProviderEnumerator = func(ctx context.Context) []builder.ProviderTarget
 
+// providerDiscoveryTimeout bounds how long the aggregate waits on ONE
+// provider's tools/list. Discovery runs in parallel (see
+// registerProviderTools), so a slow or hung provider costs at most this
+// long and never blocks the healthy ones — it just drops out of this
+// tools/list and reappears on the next one once it recovers. Kept well
+// under the providerMCPClient's 15s call timeout: discovery should be
+// fast, and a provider that can't list its tools quickly is treated as
+// absent rather than allowed to stall the whole aggregate. A var (not a
+// const) only so tests can lower it.
+var providerDiscoveryTimeout = 8 * time.Second
+
 // registerProviderTools fetches each Ready provider's tools/list and
 // registers them on srv as proxies. Errors against any one provider
 // are logged + skipped — one broken provider must not poison the
-// whole aggregate. The aggregator stays stateless: a fresh server is
-// built per request, so a provider that just became Ready shows up on
-// the very next tools/list from the client.
+// whole aggregate. Discovery is fanned out concurrently with a
+// per-provider deadline so a single slow/hung provider can neither
+// stall the others nor block the aggregate tools/list beyond
+// providerDiscoveryTimeout. The aggregator stays stateless: a fresh
+// server is built per request, so a provider that just became Ready (or
+// recovered) shows up on the very next tools/list from the client.
 func registerProviderTools(ctx context.Context, srv *mcp.Server, cfg Config) {
 	log := klogFromCfg(cfg.Deps)
 	if cfg.Providers == nil {
@@ -84,34 +99,77 @@ func registerProviderTools(ctx context.Context, srv *mcp.Server, cfg Config) {
 	// header required" — same failure mode as the UI hit before we
 	// fixed the bearer-forwarding bug.
 	cli := newProviderMCPClient(cfg.BearerToken, cfg.Cluster)
-	for _, p := range providers {
+
+	// Discover every Ready provider's tools concurrently. A sequential
+	// loop would let one slow/hung provider stall discovery of all the
+	// others (and the aggregate tools/list as a whole) for up to the
+	// client's timeout each; fanning out with a per-provider deadline
+	// caps that cost at providerDiscoveryTimeout, in parallel, and a
+	// failure just drops that one provider. Each goroutine writes only
+	// its own results[i] slot, so no mutex is needed.
+	results := make([]*providerTools, len(providers))
+	var wg sync.WaitGroup
+	for i := range providers {
+		p := providers[i]
 		if !p.Ready || p.MCPURL == "" {
 			log.Info("provider federation: skip (not ready or no MCP URL)", "provider", p.Name, "ready", p.Ready, "mcpURL", p.MCPURL)
 			continue
 		}
-		tools, err := cli.listTools(ctx, p.MCPURL)
-		if err != nil {
-			log.Info("provider federation: tools/list failed (skipping)", "provider", p.Name, "mcpURL", p.MCPURL, "err", err.Error())
+		wg.Add(1)
+		go func(i int, p builder.ProviderTarget) {
+			defer wg.Done()
+			// A panic discovering one provider must not crash the
+			// whole aggregate request.
+			defer func() {
+				if r := recover(); r != nil {
+					log.Info("provider federation: discovery panic recovered", "provider", p.Name, "panic", fmt.Sprint(r))
+				}
+			}()
+			dctx, cancel := context.WithTimeout(ctx, providerDiscoveryTimeout)
+			defer cancel()
+			tools, err := cli.listTools(dctx, p.MCPURL)
+			if err != nil {
+				log.Info("provider federation: tools/list failed (skipping)", "provider", p.Name, "mcpURL", p.MCPURL, "err", err.Error())
+				return
+			}
+			results[i] = &providerTools{provider: p, tools: tools}
+		}(i, p)
+	}
+	wg.Wait()
+
+	// Register sequentially, in the original provider order, so the
+	// aggregate tool list is deterministic across requests. AddTool on a
+	// shared *mcp.Server is not guaranteed goroutine-safe, so registration
+	// stays on this goroutine rather than inside the fan-out above.
+	for _, r := range results {
+		if r == nil {
 			continue
 		}
-		log.Info("provider federation: registering tools", "provider", p.Name, "count", len(tools))
-		for _, t := range tools {
+		log.Info("provider federation: registering tools", "provider", r.provider.Name, "count", len(r.tools))
+		for _, t := range r.tools {
 			func() {
 				// AddTool can panic on schema validation failures
 				// (missing input schema, non-object type, etc).
 				// Recover so one bad provider doesn't poison the
 				// entire aggregate's tool list.
 				defer func() {
-					if r := recover(); r != nil {
-						log.Info("provider federation: AddTool panic recovered", "provider", p.Name, "tool", t.Name, "panic", fmt.Sprint(r))
+					if rec := recover(); rec != nil {
+						log.Info("provider federation: AddTool panic recovered", "provider", r.provider.Name, "tool", t.Name, "panic", fmt.Sprint(rec))
 					}
 				}()
-				registerOneProxyTool(srv, cli, p, t)
+				registerOneProxyTool(srv, cli, r.provider, t)
 			}()
 		}
 	}
 }
 
+// providerTools is one provider's discovered tool set, carried from the
+// concurrent discovery fan-out to the sequential registration pass.
+type providerTools struct {
+	provider builder.ProviderTarget
+	tools    []discoveredTool
+}
+
 // registerOneProxyTool installs a single proxy tool on srv. Naming:
 // "<provider-slug>__<original>" so a model browsing tools/list can
 // see at a glance which provider owns which tool. The double

providers/mcp/aggregate/provider_proxy_test.go

@@ -0,0 +1,120 @@
+/*
+Copyright 2026 The Faros Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package aggregatemcp
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+
+	"github.com/faroshq/faros-kedge/pkg/virtual/builder"
+)
+
+// toolsListJSON is a minimal JSON-RPC tools/list response advertising a
+// single tool named `name`. Enough for the federation client to discover
+// and proxy it.
+func toolsListJSON(name string) string {
+	return fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"tools":[`+
+		`{"name":%q,"description":"x","inputSchema":{"type":"object"}}]}}`, name)
+}
+
+// TestRegisterProviderTools_HungProviderDoesNotBlockAggregate is the
+// hub-level resilience guard: one provider whose /mcp hangs must not take
+// the whole aggregate tools/list down with it. Discovery fans out with a
+// per-provider deadline, so the healthy provider's tools still register and
+// the hung one simply drops out of this round.
+func TestRegisterProviderTools_HungProviderDoesNotBlockAggregate(t *testing.T) {
+	// Lower the per-provider deadline so the test exercises the timeout
+	// path in milliseconds, not the production 8s.
+	prev := providerDiscoveryTimeout
+	providerDiscoveryTimeout = 300 * time.Millisecond
+	defer func() { providerDiscoveryTimeout = prev }()
+
+	healthy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(toolsListJSON("ping")))
+	}))
+	defer healthy.Close()
+
+	// The hung provider never responds within the discovery deadline: it
+	// blocks until either the federation client cancels the request (its
+	// deadline firing — the behaviour under test) or the test releases it
+	// at cleanup. release is closed before hung.Close() (defers are LIFO),
+	// so the handler returns promptly and Close doesn't stall.
+	release := make(chan struct{})
+	hung := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+		select {
+		case <-r.Context().Done():
+		case <-release:
+		}
+	}))
+	defer hung.Close()
+	defer close(release)
+
+	cfg := Config{
+		Enumerate: fakeEnumerate(nil, nil),
+		Providers: func(_ context.Context) []builder.ProviderTarget {
+			return []builder.ProviderTarget{
+				{Name: "healthy", Ready: true, MCPURL: healthy.URL},
+				{Name: "hung", Ready: true, MCPURL: hung.URL},
+			}
+		},
+	}
+
+	// newServer runs discovery synchronously; time it to prove the hung
+	// provider only cost ~one discovery deadline, not the client's 15s
+	// call timeout (and not a serial sum).
+	start := time.Now()
+	cs, cleanup := connectInMemory(t, cfg)
+	defer cleanup()
+	elapsed := time.Since(start)
+
+	if elapsed > 5*time.Second {
+		t.Fatalf("aggregate build took %s — a hung provider blocked it instead of timing out fast", elapsed)
+	}
+
+	res, err := cs.ListTools(context.Background(), nil)
+	if err != nil {
+		t.Fatalf("tools/list: %v", err)
+	}
+
+	got := map[string]bool{}
+	for _, tool := range res.Tools {
+		got[tool.Name] = true
+	}
+	if !got["healthy__ping"] {
+		t.Errorf("healthy provider's tool missing; the hung provider took the aggregate down. got tools: %v", toolNames(res.Tools))
+	}
+	for name := range got {
+		if len(name) >= 6 && name[:6] == "hung__" {
+			t.Errorf("hung provider unexpectedly contributed tool %q", name)
+		}
+	}
+}
+
+func toolNames(tools []*mcp.Tool) []string {
+	out := make([]string, 0, len(tools))
+	for _, tl := range tools {
+		out = append(out, tl.Name)
+	}
+	return out
+}