mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
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
This commit is contained in:
+81
-5
@@ -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
|
||||
|
||||
+11
-1
@@ -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
|
||||
|
||||
+226
@@ -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"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
+24
-5
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user