From 88260a386bbb18a7139d4f02368fdc0234ad860a Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 3 Sep 2026 02:45:15 +0100 Subject: [PATCH] fix(target-health): ignore context-caused federation errors mcpFederatedToolsList's fan-out and mcpFederatedToolsCall both derive their outbound context from r.Context(). A client cancel (context.Canceled) or the shared fan-out deadline expiring (context.DeadlineExceeded) is not evidence about a server's own health, so neither is recorded any more, mirroring recordTargetProxyHealth's own client-cancel rule. A server that is genuinely hung is still caught by the active probe sweep or its own per-server proxy health. Claude-Session: https://claude.ai/code/session_01EebV5jyWvYocByNQ3oss1w --- README.md | 7 +- mcp_federation.go | 23 +++++- mcp_federation_test.go | 161 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 188 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 4bd9de3..f05c4d0 100644 --- a/README.md +++ b/README.md @@ -1893,7 +1893,12 @@ plugin's config. `POST /mcp`'s `tools/list` fan-out and `tools/call` record a failure only when the outbound call itself fails (`err != nil` from the backend call) — a JSON-RPC-level `error` in the backend's own - response still counts as reachable, since the server answered. + response still counts as reachable, since the server answered. A + context-caused failure (the caller canceling, or the shared fan-out + deadline expiring) records nothing at all, the same client-cancel + rule the per-server proxy above applies; a genuinely hung server is + still caught by the active probe sweep or its own per-server proxy + health. - **Active probes are opt-in** (`targetHealth.enabled: true`) and lazy: no timer, no background goroutine running on its own. A probe sweep is triggered by `/metrics`, `GET /admin/api/targets`, `GET diff --git a/mcp_federation.go b/mcp_federation.go index 223bbd2..bdd6bd2 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -980,7 +980,20 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons // it is recorded separately, below, once the shape is known. probeStart := time.Now() resp, err := g.mcpBackendCall(fanoutCtx, targetURL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes) - g.targetHealth.record(targetKindMCP, name, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) + // F2: fanoutCtx is derived from r.Context() (above), so a + // client hang-up (context.Canceled) surfaces here as EVERY + // still-in-flight server's own error, and the shared + // toolsListBackendTimeout expiry (context.DeadlineExceeded) + // surfaces identically for a server that was genuinely slow + // AND one still parked on sem, never dialed at all — neither + // says anything about that server's own health, so neither is + // recorded, mirroring recordTargetProxyHealth's identical + // client-cancel rule (target_health.go). A server that is + // really hung is instead caught by the active probe sweep or + // its own per-request proxy health, not by this fan-out. + if err == nil || (!errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)) { + g.targetHealth.record(targetKindMCP, name, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) + } mu.Lock() defer mu.Unlock() @@ -1103,7 +1116,13 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons // not. probeStart := time.Now() resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}, mcpBackendCallResponseMaxBytes) - g.targetHealth.record(targetKindMCP, serverName, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) + // F2: same context-caused-error guard as mcpFederatedToolsList's own + // fan-out, above — ctx here derives from r.Context() too, so a client + // cancel or the shared toolsCallBackendTimeout expiry must record + // nothing rather than a false failure. + if err == nil || (!errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)) { + g.targetHealth.record(targetKindMCP, serverName, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) + } g.limiter.countTargetRequest(targetKindMCP, serverName) if err != nil { g.logf("federated tools/call: server %q: %v", serverName, err) diff --git a/mcp_federation_test.go b/mcp_federation_test.go index 94933e2..a386eb6 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -1995,3 +1995,164 @@ func TestHandleMCPFederated_ToolsCall_TargetHealth_UnreachableServer_RecordsUnhe t.Errorf("alpha state = %q, want unhealthy", got) } } + +// --- feat/target-health F2: context-caused fan-out errors record nothing --- +// +// A plain connection-refused server still recording a failure is already +// covered by TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer +// (tools/list) and TestHandleMCPFederated_ToolsCall_TargetHealth_UnreachableServer_RecordsUnhealthy +// (tools/call) above — both must keep passing unchanged by the F2 fix +// below, which only special-cases context.Canceled/context.DeadlineExceeded. + +// TestHandleMCPFederated_ToolsList_ClientCanceled_RecordsNothing proves F2: +// a client cancel mid-fan-out — context.Canceled propagating from +// r.Context() into fanoutCtx — records nothing for the servers still in +// flight, mirroring recordTargetProxyHealth's own client-cancel rule +// (target_health.go): a caller hanging up says nothing about the target's +// own health. +func TestHandleMCPFederated_ToolsList_ClientCanceled_RecordsNothing(t *testing.T) { + reached := make(chan struct{}, 2) + release := make(chan struct{}) + blocking := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + reached <- struct{}{} + <-release + }) + alpha := httptest.NewServer(blocking) + beta := httptest.NewServer(blocking) + // Registered in this order so t.Cleanup's LIFO order unblocks both + // handlers (closing release) BEFORE either server's Close is called — + // httptest.Server.Close blocks until outstanding requests complete, so + // closing it first would deadlock until the test's own timeout (F6). + t.Cleanup(alpha.Close) + t.Cleanup(beta.Close) + t.Cleanup(func() { close(release) }) + + cfg := newFederationTestConfig(alpha.URL, beta.URL, false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + + ctx, cancel := context.WithCancel(context.Background()) + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}).WithContext(ctx) + + done := make(chan struct{}) + rec := httptest.NewRecorder() + go func() { + h.ServeHTTP(rec, req) + close(done) + }() + + <-reached + <-reached + cancel() + <-done + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthUnknown { + t.Errorf("alpha state = %q, want unknown (a client cancel must record nothing)", got) + } + if got := gw.targetHealth.stateOf(targetKindMCP, "beta"); got != targetHealthUnknown { + t.Errorf("beta state = %q, want unknown (a client cancel must record nothing)", got) + } +} + +// TestHandleMCPFederated_ToolsList_SharedDeadlineExpiry_RecordsNothing +// proves F2's second case: when the fan-out's own shared context expires +// (toolsListBackendTimeout, or here the incoming request's own shorter +// deadline — see TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext's +// own doc comment for why that stands in for the real 20s budget), a +// server still parked behind the fan-out semaphore — never actually +// dialed — must not be recorded as a failure either. mcpBackendCall +// returns the identical context.DeadlineExceeded whether a server was +// slow or never contacted at all, so this case is handled by the same +// guard as the client-cancel case above, not by tracking who was dialed. +func TestHandleMCPFederated_ToolsList_SharedDeadlineExpiry_RecordsNothing(t *testing.T) { + block := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-block + })) + t.Cleanup(srv.Close) + t.Cleanup(func() { close(block) }) + + // One more server than the fan-out semaphore has slots for + // (mcpFederatedFanoutConcurrency) guarantees at least one is still + // parked on the semaphore, never dialed, when the shared deadline + // fires. + names := make([]string, 0, mcpFederatedFanoutConcurrency+1) + mcpServers := make(map[string]*TargetConfig, mcpFederatedFanoutConcurrency+1) + for i := 0; i < mcpFederatedFanoutConcurrency+1; i++ { + name := fmt.Sprintf("srv%d", i) + names = append(names, name) + mcpServers[name] = &TargetConfig{URL: srv.URL} + } + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = mcpServers + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}) + shortCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + req = req.WithContext(shortCtx) + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + for _, name := range names { + if got := gw.targetHealth.stateOf(targetKindMCP, name); got != targetHealthUnknown { + t.Errorf("mcp/%s state = %q, want unknown (a shared fan-out deadline expiry must record nothing)", name, got) + } + } +} + +// TestHandleMCPFederated_ToolsCall_ClientCanceled_RecordsNothing mirrors +// TestHandleMCPFederated_ToolsList_ClientCanceled_RecordsNothing for +// tools/call's own, separate context.WithTimeout(r.Context(), ...) call +// site. +func TestHandleMCPFederated_ToolsCall_ClientCanceled_RecordsNothing(t *testing.T) { + reached := make(chan struct{}) + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(reached) + <-release + })) + t.Cleanup(srv.Close) + t.Cleanup(func() { close(release) }) + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + + ctx, cancel := context.WithCancel(context.Background()) + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{ + JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams, + }).WithContext(ctx) + + done := make(chan struct{}) + rec := httptest.NewRecorder() + go func() { + h.ServeHTTP(rec, req) + close(done) + }() + + <-reached + cancel() + <-done + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthUnknown { + t.Errorf("alpha state = %q, want unknown (a client cancel must record nothing)", got) + } +}