From d6eebd5a84fa2f95a36bb4ad38f124eff40f1abf Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 3 Sep 2026 01:42:21 +0100 Subject: [PATCH] feat(target-health): track MCP and agent target health Add a per-replica, in-memory health tracker for MCP servers and A2A agents (target_health.go). The tracker records observations from real proxied traffic (handleTargetProxy, mcp_a2a.go) and federated calls (mcpFederatedToolsList/mcpFederatedToolsCall, mcp_federation.go), and derives a three-state readout (unknown, healthy, unhealthy) from a configurable consecutive-failure threshold. Add TargetHealthConfig to Config, validated at construction. The tracker never gates routing: federation and the target proxy keep contacting every configured server regardless of its recorded health. Extend proxyResult (routes_passthrough.go) with status and clientCanceled so handleTargetProxy can classify a proxied response without proxyUpstream knowing anything about target health. Add maybeSweepTargetHealth, an opt-in, lazy, single-flight active probe sweep with no timer and no long-lived goroutine, matching this plugin's lack of a Close hook under Traefik hot reload. A probe uses the initialize/close handshake for a Streamable HTTP MCP server, a plain unread GET for a legacy /sse server, and a GET of the agent card path for an A2A agent. handleMCPServers and handleAgents now trigger a sweep at the top of each call. Claude-Session: https://claude.ai/code/session_01EebV5jyWvYocByNQ3oss1w --- llmgateway.go | 86 ++++++- mcp_a2a.go | 12 +- mcp_a2a_test.go | 226 ++++++++++++++++++ mcp_federation.go | 12 + mcp_federation_test.go | 122 ++++++++++ routes_passthrough.go | 29 ++- target_health.go | 524 +++++++++++++++++++++++++++++++++++++++++ target_health_test.go | 413 ++++++++++++++++++++++++++++++++ 8 files changed, 1413 insertions(+), 11 deletions(-) create mode 100644 target_health.go create mode 100644 target_health_test.go diff --git a/llmgateway.go b/llmgateway.go index ba90a94..116678b 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -111,6 +111,13 @@ type Config struct { // Cache is a struct value, not a pointer, for the same reason as Retry // above: its own Enabled field is the on/off signal. Cache CacheConfig `json:"cache"` + // TargetHealth configures feat/target-health's opt-in active probe + // sweep for MCP servers and A2A agents — see TargetHealthConfig's own + // doc comment. A struct value, not a pointer, for the same + // zero-means-default reason as Breaker above: passive recording and + // exposure (GET /admin/api/targets, llmgateway_target_healthy) stay + // on unconditionally regardless of this block's presence. + TargetHealth TargetHealthConfig `json:"targetHealth"` // MaxInFlightBodyRequests caps how many unified and media JSON/ // multipart requests (routes_unified.go's runUnified via // readAndDecodeUnifiedBody; routes_media.go's decodeMediaJSONRequest/ @@ -362,6 +369,39 @@ type BreakerConfig struct { FailureThreshold int `json:"failureThreshold,omitempty"` } +// TargetHealthConfig configures feat/target-health's OPT-IN active probe +// sweep for MCP servers and A2A agents (target_health.go). Passive +// recording from real request traffic, and exposure via GET +// /admin/api/targets and the llmgateway_target_healthy gauge, are ALWAYS +// on regardless of this block's presence — Enabled here gates ONLY the +// background probe sweep (maybeSweepTargetHealth, target_health.go). The +// zero value (Enabled: false) preserves prior behavior exactly: a +// deployment that predates this feature, or one that never configures +// this block, sends no active probe traffic to any MCP server or agent. +// Field order (string before the two scalars) is fieldalignment- +// sensitive, the same convention BreakerConfig's own doc comment above +// explains. +type TargetHealthConfig struct { + // ProbeInterval is the minimum time between sweeps (a Go duration + // string, e.g. "60s"). Empty (the default) uses + // defaultTargetHealthProbeInterval (target_health.go). Must be + // between minTargetHealthProbeInterval (10s) and + // maxTargetHealthProbeInterval (1h) when set, or construction fails. + ProbeInterval string `json:"probeInterval,omitempty"` + // FailureThreshold is how many CONSECUTIVE failed observations + // (probe or traffic) mark a target unhealthy — see + // targetHealthTracker's own doc comment (target_health.go) for the + // exact state rule. 0 (the default) uses + // defaultTargetHealthFailureThreshold (3). Must be between 1 and 100, + // or construction fails. + FailureThreshold int `json:"failureThreshold,omitempty"` + // Enabled opts into the active probe sweep. false (the default) + // means no MCP server or A2A agent is ever proactively contacted for + // health purposes — passive recording from real proxied traffic + // still runs either way. + Enabled bool `json:"enabled,omitempty"` +} + // AdminConfig configures the read-only admin dashboard (spec §4, v0.2). // The zero value (Enabled: false) disables it, preserving v0.1 behavior // exactly: no /admin* route is registered, so those paths fall through @@ -561,11 +601,22 @@ type Gateway struct { // degrades to "failover never skips a candidate for request health" // rather than a nil-pointer panic. failoverHealth *requestHealthTracker - limiter *limiter - registry *modelRegistry - adapters map[string]providerAdapter - cfg *Config - auth *authStore + // targetHealth is feat/target-health's own per-pod, in-memory health + // tracker for MCP servers and A2A agents (target_health.go) — always + // constructed, here, regardless of Config.TargetHealth.Enabled: + // passive recording (handleTargetProxy, mcp_a2a.go; + // mcpFederatedToolsList/mcpFederatedToolsCall, mcp_federation.go) and + // exposure (GET /admin/api/targets, llmgateway_target_healthy) both + // stay on unconditionally, mirroring g.latency/g.provenance's own + // "always-present, usually-empty, nil-receiver-safe" convention + // below. Enabled gates only maybeSweepTargetHealth's own active + // probe sweep. + targetHealth *targetHealthTracker + limiter *limiter + registry *modelRegistry + adapters map[string]providerAdapter + cfg *Config + auth *authStore // cache is nil whenever response caching is not configured or not // usable (cfg.Cache.Enabled is false, or true with no config.Redis — // see buildResponseCache, cache.go). Every call site checks for nil @@ -624,6 +675,10 @@ type Gateway struct { // failover is Config.Failover, validated and resolved once by // newGateway (validateFailoverConfig, failover.go). failover failoverConfig + // targetHealthCfg is Config.TargetHealth, validated and resolved once + // by validateTargetHealthConfig (target_health.go) — maybeSweepTargetHealth + // reads probeInterval/enabled from this, never the raw Config. + targetHealthCfg targetHealthConfig // targetTimeout is targetClient's own resolved request timeout — // Config.RequestTimeout only; an MCP/A2A target has no per-target // override the way a provider does. proxyUpstream's own callers @@ -631,6 +686,15 @@ type Gateway struct { // idle-progress body watchdog (timeout.go) matches the // ResponseHeaderTimeout already set on targetClient's Transport. targetTimeout time.Duration + // targetHealthLastSweepUnixNano/targetHealthSweeping back + // maybeSweepTargetHealth's own single-flight, at-most-one-sweep-per- + // interval gate (target_health.go): a UnixNano timestamp and a CAS + // flag, both atomic rather than mutex-guarded, since every caller + // (serveMetrics, serveAdminTargets, handleMCPServers, handleAgents) + // reaches maybeSweepTargetHealth on a live request path and must + // never block behind a lock another such call already holds. + targetHealthLastSweepUnixNano int64 + targetHealthSweeping int32 } // telemetryStartupOnce keeps the anonymous "plugin loaded" ping to one per @@ -749,6 +813,18 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str } g.failover = failoverCfg g.failoverHealth = newRequestHealthTracker() + + // feat/target-health: validated once, here — passive recording and + // exposure stay always-on regardless of this block; Enabled gates + // only maybeSweepTargetHealth's own active probe sweep + // (target_health.go). + targetHealthCfg, err := validateTargetHealthConfig(config.TargetHealth) + if err != nil { + return nil, err + } + g.targetHealthCfg = targetHealthCfg + g.targetHealth = newTargetHealthTracker(targetHealthCfg.failureThreshold) + // feat: instrument upstream latency — always constructed, regardless // of whether Config.Metrics is nil/disabled; see g.latency's own doc // comment above for why an always-present, usually-empty store is the diff --git a/mcp_a2a.go b/mcp_a2a.go index a7168c7..1fc6c64 100644 --- a/mcp_a2a.go +++ b/mcp_a2a.go @@ -6,6 +6,7 @@ import ( "net/http" "net/url" "sort" + "time" ) // Target kinds handleTargetProxy and resolveTarget dispatch on, selecting @@ -63,6 +64,7 @@ type agentListing struct { // servers, filtered by group.allowsMCP and sorted by name for a // deterministic response. func (g *Gateway) handleMCPServers(w http.ResponseWriter, grp *group) { + g.maybeSweepTargetHealth() names := make([]string, 0, len(g.cfg.MCPServers)) for name := range g.cfg.MCPServers { if grp.allowsMCP(name) { @@ -86,6 +88,7 @@ func (g *Gateway) handleMCPServers(w http.ResponseWriter, grp *group) { // proxy URL plus its configured card path, defaulting to // defaultAgentCardPath when AgentConfig.Card is empty. func (g *Gateway) handleAgents(w http.ResponseWriter, grp *group) { + g.maybeSweepTargetHealth() names := make([]string, 0, len(g.cfg.Agents)) for name := range g.cfg.Agents { if grp.allowsAgent(name) { @@ -216,7 +219,14 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u upstreamURL += "?" + r.URL.RawQuery } - g.proxyUpstream(w, r, upstreamURL, g.targetClient, nil, nil, false, kind+" target (name "+name+")", g.targetTimeout) + // feat/target-health: passive recording (target_health.go) — latency + // measured around the whole proxyUpstream call, matching what an + // active probe of the same target measures. See + // recordTargetProxyHealth's own doc comment for the exact + // success/failure rule. + start := time.Now() + result, ok := g.proxyUpstream(w, r, upstreamURL, g.targetClient, nil, nil, false, kind+" target (name "+name+")", g.targetTimeout) + g.recordTargetProxyHealth(kind, name, result, ok, time.Since(start)) } // validateTargetURLs checks every configured MCP-server and agent entry at diff --git a/mcp_a2a_test.go b/mcp_a2a_test.go index 77842d0..7394d66 100644 --- a/mcp_a2a_test.go +++ b/mcp_a2a_test.go @@ -935,3 +935,229 @@ func TestHandleTargetProxy_PerTargetCounters_AttributedAndCollisionSafe(t *testi }) } } + +// --- feat/target-health: passive recording via handleTargetProxy --- + +func newTargetHealthProxyConfig(mcpURL string) *Config { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: mcpURL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + return cfg +} + +// TestHandleTargetProxy_TargetHealth_200_RecordsHealthy proves an +// ordinary successful proxy call records a healthy, "traffic"-sourced +// observation. +func TestHandleTargetProxy_TargetHealth_200_RecordsHealthy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := newTargetHealthProxyConfig(srv.URL) + 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 := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthHealthy { + t.Errorf("state = %q, want healthy", snap.state) + } + if snap.source != targetHealthSourceTraffic { + t.Errorf("source = %q, want %q", snap.source, targetHealthSourceTraffic) + } + if snap.consecutiveFailures != 0 { + t.Errorf("consecutiveFailures = %d, want 0", snap.consecutiveFailures) + } +} + +// TestHandleTargetProxy_TargetHealth_404_RecordsHealthy proves a 4xx +// upstream response still counts as reachable — only 5xx (or an outright +// failed attempt) is a failure. +func TestHandleTargetProxy_TargetHealth_404_RecordsHealthy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + cfg := newTargetHealthProxyConfig(srv.URL) + 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 := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthHealthy { + t.Errorf("state = %q, want healthy (a 4xx still proves the target answered)", snap.state) + } +} + +// TestHandleTargetProxy_TargetHealth_500_RecordsUnhealthy proves an +// upstream 5xx is a failure, distinct from a 4xx above. +func TestHandleTargetProxy_TargetHealth_500_RecordsUnhealthy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + cfg := newTargetHealthProxyConfig(srv.URL) + cfg.TargetHealth = TargetHealthConfig{FailureThreshold: 1} + 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 := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthUnhealthy { + t.Errorf("state = %q, want unhealthy", snap.state) + } + if snap.lastError == "" { + t.Error("want a non-empty lastError naming the 500 status") + } +} + +// TestHandleTargetProxy_TargetHealth_DeadUpstream_RecordsUnhealthy +// proves an outright failed attempt (connection refused) is a failure, +// exercising proxyResult's status=0/clientCanceled=false shape. +func TestHandleTargetProxy_TargetHealth_DeadUpstream_RecordsUnhealthy(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + deadURL := srv.URL + srv.Close() + + cfg := newTargetHealthProxyConfig(deadURL) + cfg.TargetHealth = TargetHealthConfig{FailureThreshold: 1} + 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 := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthUnhealthy { + t.Errorf("state = %q, want unhealthy", snap.state) + } +} + +// TestHandleTargetProxy_TargetHealth_HangTripsWatchdog_RecordsUnhealthy +// proves a stalled upstream that trips the idle-progress watchdog +// (timeout.go) is recorded as a failure, not silently ignored — the +// hang fixture mirrors TestHandlePassthrough_StalledUpstream_AbortedAtTimeout +// (timeout_test.go). +func TestHandleTargetProxy_TargetHealth_HangTripsWatchdog_RecordsUnhealthy(t *testing.T) { + srv, release := newStallingServer(nil) + defer release() + + cfg := newTargetHealthProxyConfig(srv.URL) + cfg.RequestTimeout = "60ms" + cfg.TargetHealth = TargetHealthConfig{FailureThreshold: 1} + 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 := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + + start := time.Now() + h.ServeHTTP(rec, req) + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Errorf("target proxy took %s to abort a 60ms requestTimeout", elapsed) + } + if rec.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String()) + } + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthUnhealthy { + t.Errorf("state = %q, want unhealthy (the watchdog firing must record a failure)", snap.state) + } +} + +// TestHandleTargetProxy_TargetHealth_ClientCanceled_RecordsNothing +// proves a canceled CLIENT request records nothing at all — the target's +// own health stays "unknown", not "unhealthy", since the client hanging +// up says nothing about the target. +func TestHandleTargetProxy_TargetHealth_ClientCanceled_RecordsNothing(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(200 * time.Millisecond) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := newTargetHealthProxyConfig(srv.URL) + 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()) + go func() { + time.Sleep(30 * time.Millisecond) + cancel() + }() + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil).WithContext(ctx) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + if snap.state != targetHealthUnknown { + t.Errorf("state = %q, want unknown (a client cancel must record nothing)", snap.state) + } +} + +// TestHandleTargetProxy_TargetHealth_A2A_UsesAgentKind proves an A2A +// target proxy call records under targetKindAgent, independent of an MCP +// server sharing the same name. +func TestHandleTargetProxy_TargetHealth_A2A_UsesAgentKind(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Agents = map[string]*AgentConfig{"alpha": {URL: srv.URL}} + 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 := httptest.NewRequest(http.MethodGet, "/a2a/alpha/tasks", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + h.ServeHTTP(httptest.NewRecorder(), req) + + if gw.targetHealth.stateOf(targetKindAgent, "alpha") != targetHealthHealthy { + t.Errorf("agent/alpha state = %q, want healthy", gw.targetHealth.stateOf(targetKindAgent, "alpha")) + } + if gw.targetHealth.stateOf(targetKindMCP, "alpha") != targetHealthUnknown { + t.Errorf("mcp/alpha state = %q, want unknown (must not be conflated with the agent of the same name)", gw.targetHealth.stateOf(targetKindMCP, "alpha")) + } +} diff --git a/mcp_federation.go b/mcp_federation.go index b916bd1..e2dd4d3 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -974,7 +974,13 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons defer func() { <-sem }() targetURL := g.cfg.MCPServers[name].URL + // feat/target-health: passive recording (target_health.go) — + // err != nil is a failure; a JSON-RPC-level resp.Error is NOT + // (the server answered, it just reported its own error), so + // 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, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) mu.Lock() defer mu.Unlock() @@ -1091,7 +1097,13 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons // important finding 4): tools/call always resolves to exactly ONE // backend, so the fan-out memory-amplification argument does not // apply here — see mcpBackendCallResponseMaxBytes' own doc comment. + // feat/target-health: passive recording (target_health.go) — measured + // around the whole outbound call, same rule as tools/list's own fan- + // out above: err != nil is a failure, a JSON-RPC-level resp.Error is + // not. + probeStart := time.Now() resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}, mcpBackendCallResponseMaxBytes) + g.targetHealth.record(targetKindMCP, serverName, 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 ed0e5a4..94933e2 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -1873,3 +1873,125 @@ func TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller(t *testing.T) t.Fatalf("tools = %+v, want alpha_search despite the session-close failure", result.Tools) } } + +// --- feat/target-health: passive recording from federation traffic --- + +// TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer proves +// tools/list's own fan-out records a failure for the one server that +// could not be reached, and a healthy observation for every server that +// answered — independent of the aggregate response degrading rather +// than failing (TestHandleMCPFederated_ToolsList_UnreachableServerDegradesNotFails +// already covers the response shape; this proves the target-health side +// effect). +func TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer(t *testing.T) { + alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}}) + cfg := newFederationTestConfig(alpha.srv.URL, "http://127.0.0.1:1", false) + cfg.TargetHealth = TargetHealthConfig{FailureThreshold: 1} + 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) + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthHealthy { + t.Errorf("alpha state = %q, want healthy", got) + } + if got := gw.targetHealth.stateOf(targetKindMCP, "beta"); got != targetHealthUnhealthy { + t.Errorf("beta state = %q, want unhealthy (unreachable)", got) + } +} + +// TestHandleMCPFederated_ToolsList_TargetHealth_JSONRPCError_StillHealthy +// proves a server that ANSWERS with its own JSON-RPC-level error is +// still recorded healthy — it responded, which is what target-health +// cares about; only a transport-level failure (err != nil from +// mcpBackendCall) counts against it. +func TestHandleMCPFederated_ToolsList_TargetHealth_JSONRPCError_StillHealthy(t *testing.T) { + alpha := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "tools unavailable"}}) + })) + defer alpha.Close() + + cfg := newFederationTestConfig(alpha.URL, alpha.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) + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthHealthy { + t.Errorf("alpha state = %q, want healthy — a JSON-RPC-level error still means the server answered", got) + } +} + +// TestHandleMCPFederated_ToolsCall_TargetHealth_RecordsResolvedServerOnly +// proves tools/call records health for the ONE resolved server, never +// the other configured-but-uncontacted one. +func TestHandleMCPFederated_ToolsCall_TargetHealth_RecordsResolvedServerOnly(t *testing.T) { + alpha := newMockJSONRPCServer(t, nil) + beta := newMockJSONRPCServer(t, nil) + cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.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) + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{ + JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("5"), Params: callParams, + })) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthHealthy { + t.Errorf("alpha state = %q, want healthy", got) + } + if got := gw.targetHealth.stateOf(targetKindMCP, "beta"); got != targetHealthUnknown { + t.Errorf("beta state = %q, want unknown (never contacted)", got) + } +} + +// TestHandleMCPFederated_ToolsCall_TargetHealth_UnreachableServer_RecordsUnhealthy +// proves a resolved-but-unreachable backend records a failure even +// though the client sees a generic JSON-RPC internal error, not a raw +// transport error. +func TestHandleMCPFederated_ToolsCall_TargetHealth_UnreachableServer_RecordsUnhealthy(t *testing.T) { + cfg := newFederationTestConfig("http://127.0.0.1:1", "http://127.0.0.1:1", false) + cfg.TargetHealth = TargetHealthConfig{FailureThreshold: 1} + 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) + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{ + JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("5"), Params: callParams, + })) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + if got := gw.targetHealth.stateOf(targetKindMCP, "alpha"); got != targetHealthUnhealthy { + t.Errorf("alpha state = %q, want unhealthy", got) + } +} diff --git a/routes_passthrough.go b/routes_passthrough.go index 37fb3a0..dd99abb 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -910,8 +910,27 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u // nil unless accountJSON was true and the response's Content-Type was // application/json. type proxyResult struct { - tee *cappedAccountingBuffer + tee *cappedAccountingBuffer + // status is the upstream response's own status code, once headers + // were actually written to w — 0 when proxyUpstream returned before + // ever writing them (a build/connection failure, or a client cancel + // before the request was even sent). Additive field for + // feat/target-health's passive recording (recordTargetProxyHealth, + // target_health.go): ok=false alone conflates "the upstream never + // answered at all" with "it answered with a 5xx", which that + // feature's failure rule needs to distinguish. Field order (pointer, + // then int, then the two bools) is fieldalignment-sensitive, the same + // convention this package's other structs already follow. + status int isJSON bool + // clientCanceled is true exactly when either of proxyUpstream's own + // errors.Is(err, context.Canceled) branches fired — the CLIENT, not + // the upstream, ended the request, whether before the upstream call + // was even sent or mid-copy of an already-started response. + // feat/target-health's passive recording must record NOTHING for + // this case (handleTargetProxy's own doc comment, mcp_a2a.go): a + // client hanging up says nothing about the target's own health. + clientCanceled bool } // proxyUpstream is the shared reverse-proxy core behind both native @@ -1054,7 +1073,7 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream cancel() if errors.Is(err, context.Canceled) { g.logf("%s: client canceled request: %v", logPrefix, err) - return proxyResult{}, false + return proxyResult{clientCanceled: true}, false } g.errorf("%s: upstream connection error: %v", logPrefix, err) writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") @@ -1099,10 +1118,10 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream // — the latter falls through to the errorf below, unchanged. if errors.Is(err, context.Canceled) { g.logf("%s: client canceled request: %v", logPrefix, err) - return proxyResult{isJSON: isJSON, tee: tee}, false + return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode, clientCanceled: true}, false } g.errorf("%s: stream response body: %v", logPrefix, err) - return proxyResult{isJSON: isJSON, tee: tee}, false + return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode}, false } - return proxyResult{isJSON: isJSON, tee: tee}, true + return proxyResult{isJSON: isJSON, tee: tee, status: resp.StatusCode}, true } diff --git a/target_health.go b/target_health.go new file mode 100644 index 0000000..33cdeb7 --- /dev/null +++ b/target_health.go @@ -0,0 +1,524 @@ +package traefikllmgateway + +import ( + "context" + "errors" + "fmt" + "net/http" + "sync" + "sync/atomic" + "time" +) + +// This file implements feat/target-health: a per-replica, in-memory +// health tracker for MCP servers and A2A agents ("targets"), which have +// no health signal today the way providers do (registry.go's discovery +// circuit breaker, failover.go's requestHealthTracker). Two feeds keep +// it current: always-on PASSIVE recording from real request traffic +// (handleTargetProxy, mcp_a2a.go; mcpFederatedToolsList/ +// mcpFederatedToolsCall, mcp_federation.go) and an opt-in, lazy ACTIVE +// probe sweep (maybeSweepTargetHealth, below). Exposure is always on +// too: GET /admin/api/targets (admin.go) and the llmgateway_target_healthy +// gauge (metrics.go) both read this tracker regardless of +// Config.TargetHealth.Enabled — that flag gates only the active probe +// sweep. This tracker is OBSERVE-ONLY: nothing in mcp_a2a.go or +// mcp_federation.go ever skips a server or agent because it reports +// unhealthy — federation routing behavior is unchanged by this feature. + +// targetHealthState is the tri-state health readout stateOf/snapshot +// report for one target, mirroring the discovery breaker's healthState +// string shape (registry.go's providerState) so the admin API and +// dashboard can treat both consistently. +type targetHealthState string + +const ( + // targetHealthUnknown: the tracker has never recorded an observation + // (probe or traffic) for this (kind, name) — a target configured but + // never yet reached. + targetHealthUnknown targetHealthState = "unknown" + // targetHealthHealthy: recorded at least once, and consecutiveFailures + // has not yet reached failureThreshold. This DELIBERATELY includes a + // target whose most recent single observation was itself a failure: + // 1 or 2 recent failures after a success still read "healthy" until + // the configured threshold is actually crossed — the same + // consecutive-failure gate the discovery circuit breaker uses + // (registry.go's providerState/recordHealthLocked), reimplemented + // independently here since this tracker carries no open/half-open + // state machine of its own (see this file's own package doc comment + // for why). + targetHealthHealthy targetHealthState = "healthy" + // targetHealthUnhealthy: consecutiveFailures >= failureThreshold. + targetHealthUnhealthy targetHealthState = "unhealthy" +) + +// Source values record's callers pass — which mechanism produced the +// last observation, exposed on GET /admin/api/targets and stored per +// entry. +const ( + targetHealthSourceProbe = "probe" + targetHealthSourceTraffic = "traffic" +) + +// targetHealthMaxErrorRunes bounds targetHealthEntry.lastError: an +// upstream MCP/A2A target's own error text is not operator-controlled +// the way its URL is, so a pathological error string must not grow this +// in-memory tracker unbounded. Runes, not bytes (truncateRunes, below) +// — a byte-level cut could split a multi-byte UTF-8 sequence. +const targetHealthMaxErrorRunes = 200 + +// targetHealthKey identifies one tracked target by routing kind +// (targetKindMCP or targetKindAgent, mcp_a2a.go — NOT scopeKindAgent, +// which is the limiter/admin-API-facing "agent" spelling for a2a +// targets) and configured name. +type targetHealthKey struct { + kind string + name string +} + +// targetHealthEntry is one target's current health state, guarded by its +// owning tracker's own mu (never its own) — mirrors latencyHistogram's +// identical "guarded by the parent store's lock" convention +// (metrics.go). +type targetHealthEntry struct { + lastCheck time.Time + lastError string + source string + latency time.Duration + consecutiveFailures int + lastOK bool +} + +// targetHealthTracker is feat/target-health's own per-pod, in-memory +// health signal for every configured MCP server and A2A agent. +// Deliberately separate from requestHealthTracker (failover.go) and the +// discovery circuit breaker (registry.go): this tracker carries no +// open/half-open/backoff state machine and never gates routing — it +// exists purely for operational visibility. +// +// Every method is nil-receiver-safe, matching this package's own +// nil-means-disabled convention (requestHealthTracker's own doc +// comment, failover.go) — a Gateway assembled directly in a test, +// bypassing newGateway, degrades to "every target reads unknown" rather +// than a nil-pointer panic, even though newGateway itself always +// constructs a real one (see Gateway.targetHealth's own doc comment, +// llmgateway.go). +type targetHealthTracker struct { + entries map[targetHealthKey]*targetHealthEntry + nowFn func() time.Time + // failureThreshold is TargetHealthConfig.FailureThreshold, validated + // and defaulted once at construction (validateTargetHealthConfig) — + // copied here rather than re-read from Config on every stateOf call, + // the same resolved-value-not-raw-config shape breakerConfig/ + // failoverConfig already use. + failureThreshold int + mu sync.Mutex +} + +// newTargetHealthTracker returns an empty tracker: every target reads +// "unknown" until its first recorded observation. threshold is the +// resolved (validated, defaulted) TargetHealthConfig.FailureThreshold. +func newTargetHealthTracker(threshold int) *targetHealthTracker { + return &targetHealthTracker{entries: make(map[targetHealthKey]*targetHealthEntry), nowFn: time.Now, failureThreshold: threshold} +} + +// record accounts one observation (a real proxied request, or an active +// probe) for (kind, name): success resets consecutiveFailures to 0 and +// clears lastError; failure increments consecutiveFailures and stores +// err's message, truncated to targetHealthMaxErrorRunes (err may be nil +// even on failure — a caller with no Go error to attach, e.g. a bad +// HTTP status alone, still leaves lastError at whatever it previously +// held rather than clearing it, since there IS a new failure to report). +// source is targetHealthSourceProbe or targetHealthSourceTraffic — +// whichever mechanism produced this observation. +func (t *targetHealthTracker) record(kind, name string, ok bool, err error, latency time.Duration, source string) { + if t == nil || name == "" { + return + } + t.mu.Lock() + defer t.mu.Unlock() + key := targetHealthKey{kind: kind, name: name} + e, exists := t.entries[key] + if !exists { + e = &targetHealthEntry{} + t.entries[key] = e + } + e.lastCheck = t.nowFn() + e.lastOK = ok + e.latency = latency + e.source = source + if ok { + e.consecutiveFailures = 0 + e.lastError = "" + return + } + e.consecutiveFailures++ + if err != nil { + e.lastError = truncateRunes(err.Error(), targetHealthMaxErrorRunes) + } +} + +// stateOf reports (kind, name)'s current targetHealthState — see this +// type's own doc comment for the exact unknown/healthy/unhealthy rule. +func (t *targetHealthTracker) stateOf(kind, name string) targetHealthState { + if t == nil { + return targetHealthUnknown + } + t.mu.Lock() + defer t.mu.Unlock() + e, ok := t.entries[targetHealthKey{kind: kind, name: name}] + if !ok { + return targetHealthUnknown + } + return t.stateForLocked(e) +} + +// stateForLocked derives e's targetHealthState. Caller must hold t.mu. +// Split out of stateOf/snapshot so both apply the identical rule. +func (t *targetHealthTracker) stateForLocked(e *targetHealthEntry) targetHealthState { + if e.consecutiveFailures >= t.failureThreshold { + return targetHealthUnhealthy + } + return targetHealthHealthy +} + +// targetHealthSnapshot is one target's fully-copied health state, safe +// to read after targetHealthTracker.snapshot returns without holding its +// lock — mirrors latencySnapshot's identical shape (metrics.go). +type targetHealthSnapshot struct { + lastCheck time.Time + state targetHealthState + lastError string + source string + latency time.Duration + consecutiveFailures int +} + +// snapshot returns a copy of (kind, name)'s current health state — +// buildAdminTargets (admin.go) and writeTargetMetrics (metrics.go) both +// read through this, once per configured target, rather than exposing +// t.entries directly. A target this tracker has never observed returns +// a zero-valued snapshot with state targetHealthUnknown. +func (t *targetHealthTracker) snapshot(kind, name string) targetHealthSnapshot { + if t == nil { + return targetHealthSnapshot{state: targetHealthUnknown} + } + t.mu.Lock() + defer t.mu.Unlock() + e, ok := t.entries[targetHealthKey{kind: kind, name: name}] + if !ok { + return targetHealthSnapshot{state: targetHealthUnknown} + } + return targetHealthSnapshot{ + state: t.stateForLocked(e), + lastCheck: e.lastCheck, + lastError: e.lastError, + source: e.source, + latency: e.latency, + consecutiveFailures: e.consecutiveFailures, + } +} + +// truncateRunes bounds s to at most n runes, decoding as UTF-8 rather +// than slicing raw bytes — a byte-level cut on an arbitrary upstream +// error string could split a multi-byte rune (escapeLabelValue's own +// doc comment, metrics.go, is the identical concern applied to a +// different string). +func truncateRunes(s string, n int) string { + r := []rune(s) + if len(r) <= n { + return s + } + return string(r[:n]) +} + +// --- config: TargetHealthConfig validation --- + +// defaultTargetHealthProbeInterval/min/max bound TargetHealthConfig. +// ProbeInterval; defaultTargetHealthFailureThreshold/max bound +// TargetHealthConfig.FailureThreshold. See TargetHealthConfig's own doc +// comment (llmgateway.go) for what each governs. +const ( + defaultTargetHealthProbeInterval = 60 * time.Second + minTargetHealthProbeInterval = 10 * time.Second + maxTargetHealthProbeInterval = time.Hour + defaultTargetHealthFailureThreshold = 3 + maxTargetHealthFailureThreshold = 100 +) + +// targetHealthConfig is TargetHealthConfig, validated and defaulted once +// at construction (newGateway, llmgateway.go) — mirrors breakerConfig/ +// failoverConfig's own resolved-value-not-raw-config shape. +type targetHealthConfig struct { + probeInterval time.Duration + failureThreshold int + enabled bool +} + +// validateTargetHealthConfig validates thc and returns the resolved +// targetHealthConfig newGateway attaches to the Gateway. thc's zero +// value (no targetHealth block at all) resolves to every default below +// unchanged: enabled stays false (a deployment that predates this +// feature, or never configures this block, sends no active probe +// traffic), matching FailoverConfig's own "off unless explicitly opted +// in" precedent (failover.go). +func validateTargetHealthConfig(thc TargetHealthConfig) (targetHealthConfig, error) { + interval := defaultTargetHealthProbeInterval + if thc.ProbeInterval != "" { + d, err := time.ParseDuration(thc.ProbeInterval) + if err != nil { + return targetHealthConfig{}, fmt.Errorf("llmgateway: targetHealth.probeInterval %q is invalid: %w", thc.ProbeInterval, err) + } + if d < minTargetHealthProbeInterval || d > maxTargetHealthProbeInterval { + return targetHealthConfig{}, fmt.Errorf("llmgateway: targetHealth.probeInterval must be between %s and %s, got %q", minTargetHealthProbeInterval, maxTargetHealthProbeInterval, thc.ProbeInterval) + } + interval = d + } + + threshold := thc.FailureThreshold + if threshold == 0 { + threshold = defaultTargetHealthFailureThreshold + } + if threshold < 1 || threshold > maxTargetHealthFailureThreshold { + return targetHealthConfig{}, fmt.Errorf("llmgateway: targetHealth.failureThreshold must be between 1 and %d, got %d", maxTargetHealthFailureThreshold, threshold) + } + + return targetHealthConfig{enabled: thc.Enabled, probeInterval: interval, failureThreshold: threshold}, nil +} + +// --- active probes: opt-in, lazy, single-flight --- + +// targetHealthMaxProbeTimeout caps how long any ONE active probe may +// run, regardless of Config.RequestTimeout — a probe is a cheap, +// synchronous health check, not real traffic, and defaultRequestTimeout +// (timeout.go) is 5 minutes: without this cap, one hung MCP server or +// agent could tie up a probe goroutine for that entire duration. +const targetHealthMaxProbeTimeout = 10 * time.Second + +// targetHealthProbeTimeout resolves the timeout ONE active probe gets: +// requestTimeout (g.targetTimeout), capped at +// targetHealthMaxProbeTimeout. Written as a small helper rather than the +// Go 1.21+ builtin min: this plugin runs interpreted under Yaegi, and +// this repo's own constraints call for stdlib-only, builtin-free helpers +// where a construct's interpreter support is unverified. +func targetHealthProbeTimeout(requestTimeout time.Duration) time.Duration { + if requestTimeout <= 0 || requestTimeout > targetHealthMaxProbeTimeout { + return targetHealthMaxProbeTimeout + } + return requestTimeout +} + +// maybeSweepTargetHealth is called at the top of every route a scrape or +// a dashboard poll can reach (serveMetrics, metrics.go; serveAdminTargets/ +// handleMCPServers/handleAgents, admin.go/mcp_a2a.go) — that traffic IS +// the heartbeat: this plugin has no Close hook under Traefik's hot +// reload (Gateway.Close's own doc comment, llmgateway.go) and therefore +// starts no timer or long-lived goroutine of its own here either, +// mirroring modelRegistry.maybeRefresh's identical reasoning +// (registry.go). +// +// Returns immediately when disabled. Otherwise: g.targetHealthSweeping +// is the single-flight gate (an atomic CAS, not a mutex, since every +// caller reaches this on a live request path and must never block +// behind a lock another such call already holds) — a CAS that fails +// means either a sweep is already in flight, in which case returning is +// exactly right regardless of how long ago it started, or another +// goroutine just won the race to start this interval's sweep, in which +// case returning is exactly right too. The interval check above the CAS +// is a cheap early-out only; the CAS itself is what actually enforces +// "at most one sweep in flight, and at most one starts per interval". +func (g *Gateway) maybeSweepTargetHealth() { + if !g.targetHealthCfg.enabled { + return + } + last := atomic.LoadInt64(&g.targetHealthLastSweepUnixNano) + if last != 0 && time.Duration(time.Now().UnixNano()-last) < g.targetHealthCfg.probeInterval { + return + } + if !atomic.CompareAndSwapInt32(&g.targetHealthSweeping, 0, 1) { + return + } + atomic.StoreInt64(&g.targetHealthLastSweepUnixNano, time.Now().UnixNano()) + go g.sweepTargetHealth() //nolint:gosec // G118: self-terminating, bounded by mcpFederatedFanoutConcurrency*targetHealthMaxProbeTimeout worst case, and always clears targetHealthSweeping on return (including on panic) — see maybeRefresh's identical reasoning, registry.go, for why this plugin accepts a bounded, self-terminating background goroutine despite having no Close hook to await it +} + +// sweepTargetHealth performs one active probe sweep: every configured +// MCP server and agent, concurrently, behind a semaphore of +// mcpFederatedFanoutConcurrency (mcp_federation.go) — reusing federation's +// own fan-out concurrency constant rather than adding a second one for +// an identical "how many outbound target requests at once" question. +// Always clears g.targetHealthSweeping on return, even on panic +// (mirrors modelRegistry.refreshProvider's own deferred-recover +// reasoning, registry.go — this goroutine has no ServeHTTP caller to +// unwind into either). +func (g *Gateway) sweepTargetHealth() { + defer atomic.StoreInt32(&g.targetHealthSweeping, 0) + defer func() { + if rec := recover(); rec != nil { + g.logf("target health: sweep panicked: %v", rec) + } + }() + + timeout := targetHealthProbeTimeout(g.targetTimeout) + + var wg sync.WaitGroup + sem := make(chan struct{}, mcpFederatedFanoutConcurrency) + + for name, tc := range g.cfg.MCPServers { + if tc == nil { + continue + } + wg.Add(1) + go func(name, url string) { + defer wg.Done() + defer func() { + if rec := recover(); rec != nil { + g.logf("target health: probe for mcp %q panicked: %v", name, rec) + } + }() + sem <- struct{}{} + defer func() { <-sem }() + g.probeMCPTarget(name, url, timeout) + }(name, tc.URL) + } + for name, ac := range g.cfg.Agents { + if ac == nil { + continue + } + wg.Add(1) + go func(name, url, card string) { + defer wg.Done() + defer func() { + if rec := recover(); rec != nil { + g.logf("target health: probe for agent %q panicked: %v", name, rec) + } + }() + sem <- struct{}{} + defer func() { <-sem }() + g.probeAgentTarget(name, url, card, timeout) + }(name, ac.URL, ac.Card) + } + wg.Wait() +} + +// probeMCPTarget performs one active probe of MCP server name at +// rawURL, bounded by timeout, and records the result with source +// "probe". A legacy HTTP+SSE-transport server (isLegacySSETransportURL, +// mcp_federation.go) gets a plain GET with its body closed unread +// (probeLegacySSE, below); every other server gets the same +// initialize/close handshake mcpBackendCall's own fallback path already +// uses (mcpBackendHandshake/mcpBackendCloseSession, mcp_federation.go) +// — a JSON-RPC error answering "initialize" is a failure carrying that +// message, exactly like a transport-level error. +func (g *Gateway) probeMCPTarget(name, rawURL string, timeout time.Duration) { + start := time.Now() + + if isLegacySSETransportURL(rawURL) { + ctx, cancel := context.WithTimeout(context.Background(), timeout) + ok, err := g.probeLegacySSE(ctx, rawURL) + cancel() + g.targetHealth.record(targetKindMCP, name, ok, err, time.Since(start), targetHealthSourceProbe) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + sessionID, err := g.mcpBackendHandshake(ctx, rawURL, mcpBackendResponseMaxBytes) + cancel() + ok := err == nil + if ok && sessionID != "" { + closeCtx, closeCancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout) + g.mcpBackendCloseSession(closeCtx, rawURL, sessionID) + closeCancel() + } + g.targetHealth.record(targetKindMCP, name, ok, err, time.Since(start), targetHealthSourceProbe) +} + +// probeLegacySSE probes a legacy HTTP+SSE-transport MCP server (URL path +// ends "/sse", isLegacySSETransportURL) with a plain GET: ok iff no +// transport error and the response status is under 500. The body is +// closed IMMEDIATELY after headers, never read — an SSE endpoint streams +// indefinitely, so reading it would hang the probe for its entire +// timeout budget on every sweep. +func (g *Gateway) probeLegacySSE(ctx context.Context, rawURL string) (bool, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) //nolint:gosec // operator-configured target URL, validated at construction (validateTargetURLs, mcp_a2a.go) + if err != nil { + return false, err + } + resp, err := g.targetClient.Do(req) + if err != nil { + return false, err + } + _ = resp.Body.Close() //nolint:errcheck // deliberately unread — see this function's own doc comment + if resp.StatusCode >= http.StatusInternalServerError { + return false, fmt.Errorf("legacy sse probe: upstream returned HTTP %d", resp.StatusCode) + } + return true, nil +} + +// probeAgentTarget performs one active probe of agent name: a GET of +// baseURL+cardPath (AgentConfig.Card, or defaultAgentCardPath when empty +// — mcp_a2a.go), bounded by timeout. ok iff no transport error and the +// response status is under 500 — a 404 counts as reachable (e.g. an +// umbrella agent whose root has no card of its own), the identical +// "any status under 500 proves the target answered" rule handleTargetProxy's +// own passive recording applies (recordTargetProxyHealth, below). The +// body is closed without being read: an agent-card probe only needs the +// status. +func (g *Gateway) probeAgentTarget(name, baseURL, cardPath string, timeout time.Duration) { + if cardPath == "" { + cardPath = defaultAgentCardPath + } + start := time.Now() + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+cardPath, nil) //nolint:gosec // operator-configured target URL, validated at construction (validateTargetURLs, mcp_a2a.go) + if err != nil { + g.targetHealth.record(targetKindAgent, name, false, err, time.Since(start), targetHealthSourceProbe) + return + } + resp, err := g.targetClient.Do(req) + if err != nil { + g.targetHealth.record(targetKindAgent, name, false, err, time.Since(start), targetHealthSourceProbe) + return + } + _ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure + ok := resp.StatusCode < http.StatusInternalServerError + var probeErr error + if !ok { + probeErr = fmt.Errorf("agent probe: upstream returned HTTP %d", resp.StatusCode) + } + g.targetHealth.record(targetKindAgent, name, ok, probeErr, time.Since(start), targetHealthSourceProbe) +} + +// --- passive recording: handleTargetProxy (mcp_a2a.go) --- + +// recordTargetProxyHealth applies handleTargetProxy's own passive +// recording rule to one proxyUpstream result (routes_passthrough.go): +// failure iff the proxy attempt itself failed for a reason other than +// the client canceling, OR the upstream answered with a 5xx; otherwise +// success — any status under 500 counts as reachable, 4xx included, +// since a 4xx still proves the target itself answered. A client cancel +// (result.clientCanceled) records nothing at all: a caller hanging up +// says nothing about the target's own health. +func (g *Gateway) recordTargetProxyHealth(kind, name string, result proxyResult, ok bool, latency time.Duration) { + if result.clientCanceled { + return + } + failed := !ok || result.status >= http.StatusInternalServerError + var err error + if failed { + err = targetProxyHealthError(ok, result.status) + } + g.targetHealth.record(kind, name, !failed, err, latency, targetHealthSourceTraffic) +} + +// targetProxyHealthError builds a short, bounded description for a +// failed passive observation. proxyUpstream itself already logged the +// real underlying error; this is only what targetHealthEntry.lastError +// — exposed via GET /admin/api/targets — shows an operator. +func targetProxyHealthError(ok bool, status int) error { + if !ok { + return errors.New("target proxy: upstream request failed") + } + return fmt.Errorf("target proxy: upstream returned HTTP %d", status) +} diff --git a/target_health_test.go b/target_health_test.go new file mode 100644 index 0000000..fc6f1cc --- /dev/null +++ b/target_health_test.go @@ -0,0 +1,413 @@ +package traefikllmgateway + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- targetHealthTracker: state machine --- + +func TestTargetHealthTracker_NeverRecorded_ReadsUnknown(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + assert.Equal(t, targetHealthUnknown, tr.stateOf(targetKindMCP, "alpha")) +} + +func TestTargetHealthTracker_NilReceiver_Safe(t *testing.T) { + t.Parallel() + var tr *targetHealthTracker + assert.NotPanics(t, func() { + tr.record(targetKindMCP, "alpha", true, nil, time.Millisecond, targetHealthSourceTraffic) + }) + assert.Equal(t, targetHealthUnknown, tr.stateOf(targetKindMCP, "alpha")) + snap := tr.snapshot(targetKindMCP, "alpha") + assert.Equal(t, targetHealthUnknown, snap.state) +} + +func TestTargetHealthTracker_SingleSuccess_ReadsHealthy(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + tr.record(targetKindMCP, "alpha", true, nil, 5*time.Millisecond, targetHealthSourceTraffic) + assert.Equal(t, targetHealthHealthy, tr.stateOf(targetKindMCP, "alpha")) +} + +// TestTargetHealthTracker_BelowThreshold_StaysHealthy proves the +// documented exception: 1 or 2 recent failures after a success still +// reads "healthy" until consecutiveFailures actually reaches the +// configured threshold (targetHealthTracker's own doc comment). +func TestTargetHealthTracker_BelowThreshold_StaysHealthy(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + tr.record(targetKindMCP, "alpha", true, nil, 0, targetHealthSourceTraffic) + tr.record(targetKindMCP, "alpha", false, errors.New("boom"), 0, targetHealthSourceTraffic) + assert.Equal(t, targetHealthHealthy, tr.stateOf(targetKindMCP, "alpha"), "1 failure of threshold 3 must still read healthy") + tr.record(targetKindMCP, "alpha", false, errors.New("boom again"), 0, targetHealthSourceTraffic) + assert.Equal(t, targetHealthHealthy, tr.stateOf(targetKindMCP, "alpha"), "2 failures of threshold 3 must still read healthy") +} + +func TestTargetHealthTracker_ReachesThreshold_ReadsUnhealthy(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + for i := 0; i < 3; i++ { + tr.record(targetKindMCP, "alpha", false, errors.New("boom"), 0, targetHealthSourceTraffic) + } + assert.Equal(t, targetHealthUnhealthy, tr.stateOf(targetKindMCP, "alpha")) +} + +func TestTargetHealthTracker_SuccessAfterUnhealthy_RecoversToHealthy(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(2) + tr.record(targetKindMCP, "alpha", false, errors.New("boom"), 0, targetHealthSourceTraffic) + tr.record(targetKindMCP, "alpha", false, errors.New("boom"), 0, targetHealthSourceTraffic) + require.Equal(t, targetHealthUnhealthy, tr.stateOf(targetKindMCP, "alpha")) + + tr.record(targetKindMCP, "alpha", true, nil, 0, targetHealthSourceTraffic) + assert.Equal(t, targetHealthHealthy, tr.stateOf(targetKindMCP, "alpha")) + snap := tr.snapshot(targetKindMCP, "alpha") + assert.Equal(t, 0, snap.consecutiveFailures, "a success must reset consecutiveFailures to 0") + assert.Empty(t, snap.lastError, "a success must clear lastError") +} + +func TestTargetHealthTracker_ErrorString_BoundedLength(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + long := strings.Repeat("x", targetHealthMaxErrorRunes+500) + tr.record(targetKindMCP, "alpha", false, errors.New(long), 0, targetHealthSourceProbe) + snap := tr.snapshot(targetKindMCP, "alpha") + assert.LessOrEqual(t, len([]rune(snap.lastError)), targetHealthMaxErrorRunes) + assert.Equal(t, strings.Repeat("x", targetHealthMaxErrorRunes), snap.lastError) +} + +func TestTargetHealthTracker_ShortError_Unchanged(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + tr.record(targetKindMCP, "alpha", false, errors.New("dial tcp: connection refused"), 0, targetHealthSourceProbe) + snap := tr.snapshot(targetKindMCP, "alpha") + assert.Equal(t, "dial tcp: connection refused", snap.lastError) +} + +func TestTargetHealthTracker_KeyedByKindAndName_Independent(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(1) + tr.record(targetKindMCP, "shared", false, errors.New("boom"), 0, targetHealthSourceTraffic) + assert.Equal(t, targetHealthUnhealthy, tr.stateOf(targetKindMCP, "shared")) + assert.Equal(t, targetHealthUnknown, tr.stateOf(targetKindAgent, "shared"), "an mcp server and an agent sharing a name must track independently") +} + +func TestTargetHealthTracker_Snapshot_FieldsPopulated(t *testing.T) { + t.Parallel() + now := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + tr := newTargetHealthTracker(3) + tr.nowFn = func() time.Time { return now } + tr.record(targetKindAgent, "beta", true, nil, 42*time.Millisecond, targetHealthSourceProbe) + + snap := tr.snapshot(targetKindAgent, "beta") + assert.Equal(t, targetHealthHealthy, snap.state) + assert.Equal(t, now, snap.lastCheck) + assert.Equal(t, 42*time.Millisecond, snap.latency) + assert.Equal(t, targetHealthSourceProbe, snap.source) + assert.Equal(t, 0, snap.consecutiveFailures) +} + +func TestTargetHealthTracker_ConcurrentRecord_NoRace(t *testing.T) { + t.Parallel() + tr := newTargetHealthTracker(3) + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + tr.record(targetKindMCP, "alpha", i%2 == 0, errors.New("x"), time.Millisecond, targetHealthSourceTraffic) + }(i) + } + wg.Wait() + _ = tr.stateOf(targetKindMCP, "alpha") +} + +// --- TargetHealthConfig validation --- + +func TestValidateTargetHealthConfig_Defaults(t *testing.T) { + t.Parallel() + thc, err := validateTargetHealthConfig(TargetHealthConfig{}) + require.NoError(t, err) + assert.False(t, thc.enabled) + assert.Equal(t, defaultTargetHealthProbeInterval, thc.probeInterval) + assert.Equal(t, defaultTargetHealthFailureThreshold, thc.failureThreshold) +} + +func TestValidateTargetHealthConfig_ExplicitEnable(t *testing.T) { + t.Parallel() + thc, err := validateTargetHealthConfig(TargetHealthConfig{Enabled: true}) + require.NoError(t, err) + assert.True(t, thc.enabled) +} + +func TestValidateTargetHealthConfig_ProbeIntervalBounds(t *testing.T) { + t.Parallel() + _, err := validateTargetHealthConfig(TargetHealthConfig{ProbeInterval: "1s"}) + assert.Error(t, err, "below the 10s floor must be rejected") + + _, err = validateTargetHealthConfig(TargetHealthConfig{ProbeInterval: "2h"}) + assert.Error(t, err, "above the 1h ceiling must be rejected") + + _, err = validateTargetHealthConfig(TargetHealthConfig{ProbeInterval: "not-a-duration"}) + assert.Error(t, err) + + thc, err := validateTargetHealthConfig(TargetHealthConfig{ProbeInterval: "30s"}) + require.NoError(t, err) + assert.Equal(t, 30*time.Second, thc.probeInterval) +} + +func TestValidateTargetHealthConfig_FailureThresholdBounds(t *testing.T) { + t.Parallel() + _, err := validateTargetHealthConfig(TargetHealthConfig{FailureThreshold: -1}) + assert.Error(t, err) + + _, err = validateTargetHealthConfig(TargetHealthConfig{FailureThreshold: maxTargetHealthFailureThreshold + 1}) + assert.Error(t, err) + + thc, err := validateTargetHealthConfig(TargetHealthConfig{FailureThreshold: 10}) + require.NoError(t, err) + assert.Equal(t, 10, thc.failureThreshold) +} + +func TestNewGateway_TargetHealth_InvalidConfig_ConstructionError(t *testing.T) { + t.Parallel() + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.TargetHealth = TargetHealthConfig{ProbeInterval: "1s"} + _, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + assert.Error(t, err) +} + +func TestNewGateway_TargetHealth_DisabledByDefault_TrackerStillConstructed(t *testing.T) { + t.Parallel() + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + require.NoError(t, err) + gw, ok := h.(*Gateway) + require.True(t, ok) + require.NotNil(t, gw.targetHealth, "the tracker must always be constructed regardless of Enabled — passive recording/exposure stay on unconditionally") + assert.False(t, gw.targetHealthCfg.enabled) +} + +// --- maybeSweepTargetHealth: disabled / enabled / single-flight --- + +func newTargetHealthSweepGateway(t *testing.T, enabled bool, probeInterval time.Duration, mcpURL string) *Gateway { + t.Helper() + adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")} + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {}}, + MCPServers: map[string]*TargetConfig{"alpha": {URL: mcpURL}}, + } + reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {}) + require.NoError(t, err) + return &Gateway{ + cfg: cfg, + registry: reg, + targetHealth: newTargetHealthTracker(3), + targetHealthCfg: targetHealthConfig{enabled: enabled, probeInterval: probeInterval, failureThreshold: 3}, + targetClient: http.DefaultClient, + targetTimeout: 5 * time.Second, + name: "test", + } +} + +func TestMaybeSweepTargetHealth_Disabled_NoOutboundRequests(t *testing.T) { + t.Parallel() + hits := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits++ + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + gw := newTargetHealthSweepGateway(t, false, 10*time.Millisecond, srv.URL) + gw.maybeSweepTargetHealth() + // maybeSweepTargetHealth spawns a goroutine when it decides to sweep; + // since it must return immediately when disabled, there is nothing to + // wait on — a brief sleep would only prove absence of a signal we + // never expect to see. + time.Sleep(20 * time.Millisecond) + assert.Equal(t, 0, hits, "a disabled targetHealth must never send an outbound probe") + assert.Equal(t, targetHealthUnknown, gw.targetHealth.stateOf(targetKindMCP, "alpha")) +} + +func TestMaybeSweepTargetHealth_Enabled_ProbesConfiguredServer(t *testing.T) { + t.Parallel() + var hits int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","result":{"protocolVersion":"2025-11-25","capabilities":{}}}`)) + })) + defer srv.Close() + + gw := newTargetHealthSweepGateway(t, true, time.Hour, srv.URL) + gw.maybeSweepTargetHealth() + waitForCondition(t, func() bool { return atomic.LoadInt32(&hits) > 0 }, time.Second) + waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindMCP, "alpha") == targetHealthHealthy }, time.Second) + + snap := gw.targetHealth.snapshot(targetKindMCP, "alpha") + assert.Equal(t, targetHealthSourceProbe, snap.source) +} + +func TestMaybeSweepTargetHealth_ConcurrentTriggers_SingleFlight(t *testing.T) { + t.Parallel() + var hits int32 + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + <-release + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","result":{}}`)) + })) + defer srv.Close() + + gw := newTargetHealthSweepGateway(t, true, time.Hour, srv.URL) + + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + gw.maybeSweepTargetHealth() + }() + } + wg.Wait() + + // Give the single in-flight probe time to actually reach the server, + // then release it — regardless of how many goroutines called + // maybeSweepTargetHealth concurrently, at most one outbound request + // must have been sent. + waitForCondition(t, func() bool { return atomic.LoadInt32(&hits) > 0 }, time.Second) + close(release) + waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindMCP, "alpha") != targetHealthUnknown }, time.Second) + assert.Equal(t, int32(1), atomic.LoadInt32(&hits), "concurrent triggers within one interval must produce exactly one sweep") +} + +// TestMaybeSweepTargetHealth_LegacySSE_ClosesBodyWithoutReading proves a +// legacy "/sse" MCP server's probe uses a plain GET and closes the body +// immediately: a test server that never closes its stream must not hang +// the probe. +func TestMaybeSweepTargetHealth_LegacySSE_ClosesBodyWithoutReading(t *testing.T) { + t.Parallel() + headersSent := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + close(headersSent) + flusher, ok := w.(http.Flusher) + if ok { + flusher.Flush() + } + <-r.Context().Done() // never closes on its own; relies on the probe not reading + })) + defer srv.Close() + + adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")} + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {}}, + MCPServers: map[string]*TargetConfig{"alpha": {URL: srv.URL + "/sse"}}, + } + reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {}) + require.NoError(t, err) + gw := &Gateway{ + cfg: cfg, + registry: reg, + targetHealth: newTargetHealthTracker(3), + targetHealthCfg: targetHealthConfig{enabled: true, probeInterval: time.Hour, failureThreshold: 3}, + targetClient: http.DefaultClient, + targetTimeout: 5 * time.Second, + name: "test", + } + gw.maybeSweepTargetHealth() + waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindMCP, "alpha") != targetHealthUnknown }, 2*time.Second) + assert.Equal(t, targetHealthHealthy, gw.targetHealth.stateOf(targetKindMCP, "alpha")) +} + +func TestMaybeSweepTargetHealth_AgentReturns404_ReadsHealthy(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")} + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {}}, + Agents: map[string]*AgentConfig{"beta": {URL: srv.URL}}, + } + reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {}) + require.NoError(t, err) + gw := &Gateway{ + cfg: cfg, + registry: reg, + targetHealth: newTargetHealthTracker(3), + targetHealthCfg: targetHealthConfig{enabled: true, probeInterval: time.Hour, failureThreshold: 3}, + targetClient: http.DefaultClient, + targetTimeout: 5 * time.Second, + name: "test", + } + gw.maybeSweepTargetHealth() + waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindAgent, "beta") != targetHealthUnknown }, time.Second) + assert.Equal(t, targetHealthHealthy, gw.targetHealth.stateOf(targetKindAgent, "beta"), "a 404 agent-card fetch counts as reachable") +} + +func TestMaybeSweepTargetHealth_AgentReturns503_RecordsFailure(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer srv.Close() + + adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")} + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {}}, + Agents: map[string]*AgentConfig{"beta": {URL: srv.URL}}, + } + reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {}) + require.NoError(t, err) + gw := &Gateway{ + cfg: cfg, + registry: reg, + targetHealth: newTargetHealthTracker(1), + targetHealthCfg: targetHealthConfig{enabled: true, probeInterval: time.Hour, failureThreshold: 1}, + targetClient: http.DefaultClient, + targetTimeout: 5 * time.Second, + name: "test", + } + gw.maybeSweepTargetHealth() + waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindAgent, "beta") != targetHealthUnknown }, time.Second) + assert.Equal(t, targetHealthUnhealthy, gw.targetHealth.stateOf(targetKindAgent, "beta")) +} + +// waitForCondition polls cond until it returns true or deadline elapses, +// failing the test on timeout — used throughout this file's sweep tests +// since the sweep itself runs in a background goroutine with no +// synchronous completion signal (maybeSweepTargetHealth's own doc +// comment: no timers, no long-lived goroutines, fire-and-forget). +func waitForCondition(t *testing.T, cond func() bool, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(2 * time.Millisecond) + } + if !cond() { + t.Fatalf("condition not met within %s", timeout) + } +}