Skip to content

Commit c38d93b

Browse files
mjudeikisclaude
mjudeikis
and
claude
authored
fix(mcp): bound + parallelize provider discovery so one bad provider can't stall the aggregate (#373)
The hub-side federation already skipped providers that errored (EOF/405) and recovered AddTool panics, so a crashing provider couldn't poison the aggregate tool list. The remaining gap was a SLOW/HUNG provider: discovery was a sequential loop with the client's 15s per-request timeout, so one hung provider stalled the whole tools/list — and serialized every healthy provider behind it. Fan out per-provider tools/list concurrently, each under a short providerDiscoveryTimeout (8s, well under the 15s call timeout). A hung provider now costs at most that one deadline, in parallel, then drops out of this round and reappears on the next once it recovers. Tools are still registered sequentially in the original provider order (AddTool isn't guaranteed goroutine-safe) so the aggregate tool list stays deterministic. Discovery goroutines also recover panics individually. Adds a regression test: a healthy provider + a provider whose /mcp hangs; asserts the healthy provider's tool still registers, the hung one drops out, and the aggregate builds in well under the client timeout. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 64db37d commit c38d93b

2 files changed

Lines changed: 190 additions & 12 deletions

File tree

‎providers/mcp/aggregate/provider_proxy.go‎

Lines changed: 70 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
"io"
4545
"net/http"
4646
"strings"
47+
"sync"
4748
"time"
4849

4950
"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -57,12 +58,26 @@ import (
5758
// aggregator stays decoupled from pkg/hub/providers.
5859
type ProviderEnumerator = func(ctx context.Context) []builder.ProviderTarget
5960

61+
// providerDiscoveryTimeout bounds how long the aggregate waits on ONE
62+
// provider's tools/list. Discovery runs in parallel (see
63+
// registerProviderTools), so a slow or hung provider costs at most this
64+
// long and never blocks the healthy ones — it just drops out of this
65+
// tools/list and reappears on the next one once it recovers. Kept well
66+
// under the providerMCPClient's 15s call timeout: discovery should be
67+
// fast, and a provider that can't list its tools quickly is treated as
68+
// absent rather than allowed to stall the whole aggregate. A var (not a
69+
// const) only so tests can lower it.
70+
var providerDiscoveryTimeout = 8 * time.Second
71+
6072
// registerProviderTools fetches each Ready provider's tools/list and
6173
// registers them on srv as proxies. Errors against any one provider
6274
// are logged + skipped — one broken provider must not poison the
63-
// whole aggregate. The aggregator stays stateless: a fresh server is
64-
// built per request, so a provider that just became Ready shows up on
65-
// the very next tools/list from the client.
75+
// whole aggregate. Discovery is fanned out concurrently with a
76+
// per-provider deadline so a single slow/hung provider can neither
77+
// stall the others nor block the aggregate tools/list beyond
78+
// providerDiscoveryTimeout. The aggregator stays stateless: a fresh
79+
// server is built per request, so a provider that just became Ready (or
80+
// recovered) shows up on the very next tools/list from the client.
6681
func registerProviderTools(ctx context.Context, srv *mcp.Server, cfg Config) {
6782
log := klogFromCfg(cfg.Deps)
6883
if cfg.Providers == nil {
@@ -84,34 +99,77 @@ func registerProviderTools(ctx context.Context, srv *mcp.Server, cfg Config) {
8499
// header required" — same failure mode as the UI hit before we
85100
// fixed the bearer-forwarding bug.
86101
cli := newProviderMCPClient(cfg.BearerToken, cfg.Cluster)
87-
for _, p := range providers {
102+
103+
// Discover every Ready provider's tools concurrently. A sequential
104+
// loop would let one slow/hung provider stall discovery of all the
105+
// others (and the aggregate tools/list as a whole) for up to the
106+
// client's timeout each; fanning out with a per-provider deadline
107+
// caps that cost at providerDiscoveryTimeout, in parallel, and a
108+
// failure just drops that one provider. Each goroutine writes only
109+
// its own results[i] slot, so no mutex is needed.
110+
results := make([]*providerTools, len(providers))
111+
var wg sync.WaitGroup
112+
for i := range providers {
113+
p := providers[i]
88114
if !p.Ready || p.MCPURL == "" {
89115
log.Info("provider federation: skip (not ready or no MCP URL)", "provider", p.Name, "ready", p.Ready, "mcpURL", p.MCPURL)
90116
continue
91117
}
92-
tools, err := cli.listTools(ctx, p.MCPURL)
93-
if err != nil {
94-
log.Info("provider federation: tools/list failed (skipping)", "provider", p.Name, "mcpURL", p.MCPURL, "err", err.Error())
118+
wg.Add(1)
119+
go func(i int, p builder.ProviderTarget) {
120+
defer wg.Done()
121+
// A panic discovering one provider must not crash the
122+
// whole aggregate request.
123+
defer func() {
124+
if r := recover(); r != nil {
125+
log.Info("provider federation: discovery panic recovered", "provider", p.Name, "panic", fmt.Sprint(r))
126+
}
127+
}()
128+
dctx, cancel := context.WithTimeout(ctx, providerDiscoveryTimeout)
129+
defer cancel()
130+
tools, err := cli.listTools(dctx, p.MCPURL)
131+
if err != nil {
132+
log.Info("provider federation: tools/list failed (skipping)", "provider", p.Name, "mcpURL", p.MCPURL, "err", err.Error())
133+
return
134+
}
135+
results[i] = &providerTools{provider: p, tools: tools}
136+
}(i, p)
137+
}
138+
wg.Wait()
139+
140+
// Register sequentially, in the original provider order, so the
141+
// aggregate tool list is deterministic across requests. AddTool on a
142+
// shared *mcp.Server is not guaranteed goroutine-safe, so registration
143+
// stays on this goroutine rather than inside the fan-out above.
144+
for _, r := range results {
145+
if r == nil {
95146
continue
96147
}
97-
log.Info("provider federation: registering tools", "provider", p.Name, "count", len(tools))
98-
for _, t := range tools {
148+
log.Info("provider federation: registering tools", "provider", r.provider.Name, "count", len(r.tools))
149+
for _, t := range r.tools {
99150
func() {
100151
// AddTool can panic on schema validation failures
101152
// (missing input schema, non-object type, etc).
102153
// Recover so one bad provider doesn't poison the
103154
// entire aggregate's tool list.
104155
defer func() {
105-
if r := recover(); r != nil {
106-
log.Info("provider federation: AddTool panic recovered", "provider", p.Name, "tool", t.Name, "panic", fmt.Sprint(r))
156+
if rec := recover(); rec != nil {
157+
log.Info("provider federation: AddTool panic recovered", "provider", r.provider.Name, "tool", t.Name, "panic", fmt.Sprint(rec))
107158
}
108159
}()
109-
registerOneProxyTool(srv, cli, p, t)
160+
registerOneProxyTool(srv, cli, r.provider, t)
110161
}()
111162
}
112163
}
113164
}
114165

166+
// providerTools is one provider's discovered tool set, carried from the
167+
// concurrent discovery fan-out to the sequential registration pass.
168+
type providerTools struct {
169+
provider builder.ProviderTarget
170+
tools []discoveredTool
171+
}
172+
115173
// registerOneProxyTool installs a single proxy tool on srv. Naming:
116174
// "<provider-slug>__<original>" so a model browsing tools/list can
117175
// see at a glance which provider owns which tool. The double

‎providers/mcp/aggregate/provider_proxy_test.go‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*
2+
Copyright 2026 The Faros Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package aggregatemcp
18+
19+
import (
20+
"context"
21+
"fmt"
22+
"net/http"
23+
"net/http/httptest"
24+
"testing"
25+
"time"
26+
27+
"github.com/modelcontextprotocol/go-sdk/mcp"
28+
29+
"github.com/faroshq/faros-kedge/pkg/virtual/builder"
30+
)
31+
32+
// toolsListJSON is a minimal JSON-RPC tools/list response advertising a
33+
// single tool named `name`. Enough for the federation client to discover
34+
// and proxy it.
35+
func toolsListJSON(name string) string {
36+
return fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":{"tools":[`+
37+
`{"name":%q,"description":"x","inputSchema":{"type":"object"}}]}}`, name)
38+
}
39+
40+
// TestRegisterProviderTools_HungProviderDoesNotBlockAggregate is the
41+
// hub-level resilience guard: one provider whose /mcp hangs must not take
42+
// the whole aggregate tools/list down with it. Discovery fans out with a
43+
// per-provider deadline, so the healthy provider's tools still register and
44+
// the hung one simply drops out of this round.
45+
func TestRegisterProviderTools_HungProviderDoesNotBlockAggregate(t *testing.T) {
46+
// Lower the per-provider deadline so the test exercises the timeout
47+
// path in milliseconds, not the production 8s.
48+
prev := providerDiscoveryTimeout
49+
providerDiscoveryTimeout = 300 * time.Millisecond
50+
defer func() { providerDiscoveryTimeout = prev }()
51+
52+
healthy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
53+
w.Header().Set("Content-Type", "application/json")
54+
_, _ = w.Write([]byte(toolsListJSON("ping")))
55+
}))
56+
defer healthy.Close()
57+
58+
// The hung provider never responds within the discovery deadline: it
59+
// blocks until either the federation client cancels the request (its
60+
// deadline firing — the behaviour under test) or the test releases it
61+
// at cleanup. release is closed before hung.Close() (defers are LIFO),
62+
// so the handler returns promptly and Close doesn't stall.
63+
release := make(chan struct{})
64+
hung := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
65+
select {
66+
case <-r.Context().Done():
67+
case <-release:
68+
}
69+
}))
70+
defer hung.Close()
71+
defer close(release)
72+
73+
cfg := Config{
74+
Enumerate: fakeEnumerate(nil, nil),
75+
Providers: func(_ context.Context) []builder.ProviderTarget {
76+
return []builder.ProviderTarget{
77+
{Name: "healthy", Ready: true, MCPURL: healthy.URL},
78+
{Name: "hung", Ready: true, MCPURL: hung.URL},
79+
}
80+
},
81+
}
82+
83+
// newServer runs discovery synchronously; time it to prove the hung
84+
// provider only cost ~one discovery deadline, not the client's 15s
85+
// call timeout (and not a serial sum).
86+
start := time.Now()
87+
cs, cleanup := connectInMemory(t, cfg)
88+
defer cleanup()
89+
elapsed := time.Since(start)
90+
91+
if elapsed > 5*time.Second {
92+
t.Fatalf("aggregate build took %s — a hung provider blocked it instead of timing out fast", elapsed)
93+
}
94+
95+
res, err := cs.ListTools(context.Background(), nil)
96+
if err != nil {
97+
t.Fatalf("tools/list: %v", err)
98+
}
99+
100+
got := map[string]bool{}
101+
for _, tool := range res.Tools {
102+
got[tool.Name] = true
103+
}
104+
if !got["healthy__ping"] {
105+
t.Errorf("healthy provider's tool missing; the hung provider took the aggregate down. got tools: %v", toolNames(res.Tools))
106+
}
107+
for name := range got {
108+
if len(name) >= 6 && name[:6] == "hung__" {
109+
t.Errorf("hung provider unexpectedly contributed tool %q", name)
110+
}
111+
}
112+
}
113+
114+
func toolNames(tools []*mcp.Tool) []string {
115+
out := make([]string, 0, len(tools))
116+
for _, tl := range tools {
117+
out = append(out, tl.Name)
118+
}
119+
return out
120+
}

0 commit comments

Comments
 (0)