traefik-llmgateway/target_health_test.go

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"))
}

// TestTargetHealthTracker_FirstObservationFailure_ReadsUnknown proves F4:
// a target whose only observations so far are failures below
// failureThreshold reads "unknown", not "healthy" — the tracker has
// never actually seen it succeed, so "healthy" would be a false claim
// even though the failure-threshold gate has not fired. The threshold
// still protects against a transient first failure: this stays
// "unknown", it does not jump straight to "unhealthy".
func TestTargetHealthTracker_FirstObservationFailure_ReadsUnknown(t *testing.T) {
	t.Parallel()
	tr := newTargetHealthTracker(3)
	tr.record(targetKindMCP, "alpha", "", false, errors.New("boom"), 0, targetHealthSourceTraffic)
	assert.Equal(t, targetHealthUnknown, tr.stateOf(targetKindMCP, "alpha"), "never having succeeded must read unknown, not healthy, below threshold")

	tr.record(targetKindMCP, "alpha", "", false, errors.New("boom again"), 0, targetHealthSourceTraffic)
	assert.Equal(t, targetHealthUnknown, tr.stateOf(targetKindMCP, "alpha"), "still unknown after a second failure, still below threshold")
}

// TestTargetHealthTracker_SuccessThenFewFailures_ReadsHealthy re-proves
// TestTargetHealthTracker_BelowThreshold_StaysHealthy's claim in F4
// terms: once a target HAS succeeded at least once, 1-2 subsequent
// failures below threshold still read "healthy" — F4 only changes the
// NEVER-succeeded case.
func TestTargetHealthTracker_SuccessThenFewFailures_ReadsHealthy(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)
	tr.record(targetKindMCP, "alpha", "", false, errors.New("boom again"), 0, targetHealthSourceTraffic)
	assert.Equal(t, targetHealthHealthy, 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)
}

// TestTargetHealthTracker_Record_ScrubsCredentialsFromLastError proves F1:
// record scrubs an error's embedded target URL through sanitizeProviderErr
// before storing lastError, so a target URL configured with an embedded
// credential (an "api-key" query parameter, or userinfo) never survives
// into GET /admin/api/targets.
func TestTargetHealthTracker_Record_ScrubsCredentialsFromLastError(t *testing.T) {
	t.Parallel()

	t.Run("query parameter credential", func(t *testing.T) {
		t.Parallel()
		tr := newTargetHealthTracker(3)
		rawURL := "https://svc.internal/mcp?api-key=SECRETVALUE"
		err := errors.New(`Get "` + rawURL + `": dial tcp: connection refused`)
		tr.record(targetKindMCP, "alpha", rawURL, false, err, 0, targetHealthSourceProbe)
		snap := tr.snapshot(targetKindMCP, "alpha")
		assert.NotContains(t, snap.lastError, "SECRETVALUE")
		assert.NotContains(t, snap.lastError, "api-key=")
	})

	t.Run("userinfo credential", func(t *testing.T) {
		t.Parallel()
		tr := newTargetHealthTracker(3)
		rawURL := "https://[email protected]/mcp"
		err := errors.New(`Get "` + rawURL + `": dial tcp: connection refused`)
		tr.record(targetKindMCP, "alpha", rawURL, false, err, 0, targetHealthSourceProbe)
		snap := tr.snapshot(targetKindMCP, "alpha")
		assert.NotContains(t, snap.lastError, "tok3n")
	})
}

// TestTargetHealthTracker_Record_ScrubsBeforeTruncate proves F1: the URL
// scrub runs against err.Error() BEFORE truncateRunes bounds it to
// targetHealthMaxErrorRunes — even when the embedded URL straddles the
// truncation point, no fragment of the secret survives. G2 (feat/target-
// health review round 2): the prefix is short enough that the SECRET
// itself — not merely the surrounding URL — straddles the 200-rune cut,
// so this test actually distinguishes scrub-before-truncate from
// truncate-before-scrub; a longer prefix that pushes the whole secret
// past the cut point would pass under either ordering, proving nothing.
func TestTargetHealthTracker_Record_ScrubsBeforeTruncate(t *testing.T) {
	t.Parallel()
	tr := newTargetHealthTracker(3)
	rawURL := "https://svc.internal/mcp?api-key=" + strings.Repeat("S", 60)
	prefix := strings.Repeat("p", 150)
	require.Less(t, len([]rune(prefix))+strings.Index(rawURL, "api-key="), targetHealthMaxErrorRunes, "the secret must start before the truncation point")
	require.Greater(t, len([]rune(prefix))+len([]rune(rawURL)), targetHealthMaxErrorRunes, "the url must straddle the truncation point")
	errText := prefix + " " + rawURL + " connection refused"

	tr.record(targetKindMCP, "alpha", rawURL, false, errors.New(errText), 0, targetHealthSourceProbe)
	snap := tr.snapshot(targetKindMCP, "alpha")
	assert.LessOrEqual(t, len([]rune(snap.lastError)), targetHealthMaxErrorRunes)
	assert.NotContains(t, snap.lastError, "SSSS", "no fragment of the secret must survive truncation")
	assert.NotContains(t, snap.lastError, "api-key=")
}

// TestTargetHealthTracker_Record_ScrubsPasswordMaskedURLForm proves G1
// (feat/target-health review round 2): net/http's own *url.Error masks a
// URL's password ("user:pass@" -> "user:***@", and "user:" with an empty
// password the same way — stripPassword, net/http/client.go) BEFORE the
// error text is ever built, whenever the dialed URL's userinfo carries a
// password. sanitizeProviderErr (admin.go) only matches rawURL verbatim,
// so for such a URL it finds nothing to scrub — leaking the password AND
// any query-string credential riding the same URL (an "api-key" value).
// record must scrub against BOTH the raw URL and net/http's own masked
// form. Every case uses a REAL error from http.Client.Do against
// 127.0.0.1:1 (this repo's own "reserved, never listening" convention —
// see provider_openai_test.go), never a hand-written string: the bug is
// about the exact shape net/http actually produces.
func TestTargetHealthTracker_Record_ScrubsPasswordMaskedURLForm(t *testing.T) {
	t.Parallel()
	// secretPart is the synthetic password these fixtures plant in a target
	// URL, and credFixtureURL assembles the URL around it at run time. The
	// parts stay separate in the source deliberately: a whole URL literal
	// carrying userinfo credentials trips this repository's own pre-commit
	// secret scan, and these fixtures exist to prove exactly such a URL
	// never reaches lastError.
	const secretPart = "sup3rsecret"
	credFixtureURL := func(userinfo, hostPath string) string {
		return "http://" + userinfo + "@" + hostPath
	}
	cases := []struct {
		name        string
		rawURL      string
		mustNotHave []string
	}{
		{
			name:        "userinfo with password",
			rawURL:      credFixtureURL("user:"+secretPart, "127.0.0.1:1/mcp"),
			mustNotHave: []string{"user:" + secretPart + "@", "user:***@", secretPart},
		},
		{
			name:        "token as username with empty password separator",
			rawURL:      credFixtureURL("TOKENASUSER:", "127.0.0.1:1/mcp"),
			mustNotHave: []string{"TOKENASUSER:@", "TOKENASUSER:***@", "TOKENASUSER"},
		},
		{
			name:        "userinfo with password and api-key query on the same URL",
			rawURL:      credFixtureURL("user:"+secretPart, "127.0.0.1:1/mcp?api-key=SECRETVALUE"),
			mustNotHave: []string{"user:" + secretPart + "@", "user:***@", secretPart, "SECRETVALUE", "api-key="},
		},
		{
			// Existing case (F1): userinfo with NO password has nothing for
			// net/http to mask, so the raw URL survives into the error
			// verbatim — must keep working unchanged by the G1 fix.
			name:        "token as username, no password separator",
			rawURL:      "http://[email protected]:1/mcp",
			mustNotHave: []string{"tok3n"},
		},
		{
			// Existing case (F1): no userinfo at all, so there is nothing
			// for net/http to mask — must keep working unchanged.
			name:        "bare api-key query, no userinfo",
			rawURL:      "http://127.0.0.1:1/mcp?api-key=SECRETVALUE",
			mustNotHave: []string{"SECRETVALUE", "api-key="},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			req, err := http.NewRequest(http.MethodGet, tc.rawURL, nil)
			require.NoError(t, err)
			_, doErr := http.DefaultClient.Do(req) //nolint:bodyclose // Do returns a nil body alongside a non-nil error
			require.Error(t, doErr, "127.0.0.1:1 must refuse the connection")

			tr := newTargetHealthTracker(3)
			tr.record(targetKindMCP, "alpha", tc.rawURL, false, doErr, 0, targetHealthSourceProbe)
			snap := tr.snapshot(targetKindMCP, "alpha")
			for _, secret := range tc.mustNotHave {
				assert.NotContains(t, snap.lastError, secret)
			}
		})
	}
}

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()
	// F6, feat/target-health review: atomic even though a disabled sweep
	// must never actually invoke this handler — a regression that DID
	// invoke it would otherwise race this counter's read below under
	// -race, obscuring the real bug (an unexpected outbound probe) behind
	// a data-race report instead of the intended assert.Equal failure.
	var hits int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.AddInt32(&hits, 1)
		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, int32(0), atomic.LoadInt32(&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":{}}`))
	}))
	// F6, feat/target-health review: registered in this order so
	// t.Cleanup's LIFO order unblocks the handler (closing release)
	// BEFORE srv.Close runs — httptest.Server.Close blocks until
	// outstanding requests complete, and the assertions below can
	// t.Fatalf/t.Errorf before ever reaching the explicit close(release)
	// further down; without this a failing assertion left the handler
	// permanently blocked and srv.Close hung until the whole test
	// binary's 10-minute timeout, rather than failing fast. unblock is
	// idempotent (sync.Once) since the happy path below also calls it
	// directly, before waiting on the sweep's resulting state.
	var unblockOnce sync.Once
	unblock := func() { unblockOnce.Do(func() { close(release) }) }
	t.Cleanup(srv.Close)
	t.Cleanup(unblock)

	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)
	unblock()
	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_IntervalEnforcedAcrossRapidCalls proves F3:
// the interval load (top of maybeSweepTargetHealth) and the CAS that wins
// the single-flight flag are not atomic with each other. Without a
// re-check under the CAS, a caller whose own interval read predates an
// entirely separate sweep's start-and-finish could still win the CAS
// once that sweep clears the flag, and spawn a second sweep inside the
// same interval — worst at startup, when every caller's first read sees
// last == 0. A fast-completing first sweep, immediately followed by a
// burst of concurrent callers, must still yield exactly one sweep within
// probeInterval; a call after the interval has genuinely elapsed must
// start a second one. G5, feat/target-health review round 2: this test
// proves single-flight ORDERING, not real elapsed time, so it must not
// depend on wall-clock sleeps at all — probeInterval is a full hour (so
// the burst phase's real, sub-millisecond elapsed time can never
// accidentally read as elapsed) and the "interval has passed" half is
// forced directly via targetHealthLastSweepUnixNano, never by sleeping
// for probeInterval.
func TestMaybeSweepTargetHealth_IntervalEnforcedAcrossRapidCalls(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":{}}`))
	}))
	defer srv.Close()

	gw := newTargetHealthSweepGateway(t, true, time.Hour, srv.URL)

	// First sweep: the probe answers immediately, so this completes fast
	// and clears targetHealthSweeping well inside probeInterval.
	gw.maybeSweepTargetHealth()
	waitForCondition(t, func() bool { return atomic.LoadInt32(&hits) > 0 }, time.Second)
	waitForCondition(t, func() bool { return gw.targetHealth.stateOf(targetKindMCP, "alpha") != targetHealthUnknown }, time.Second)

	// A burst of concurrent callers immediately after — real elapsed time
	// is milliseconds against a one-hour probeInterval — must not start a
	// second sweep.
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			gw.maybeSweepTargetHealth()
		}()
	}
	wg.Wait()
	if got := atomic.LoadInt32(&hits); got != 1 {
		t.Errorf("hits = %d, want exactly 1 (a burst of concurrent callers inside probeInterval must not start a second sweep)", got)
	}

	// Force the interval to read as elapsed directly, rather than
	// sleeping for real: a call after that must start a second sweep.
	atomic.StoreInt64(&gw.targetHealthLastSweepUnixNano, time.Now().Add(-2*time.Hour).UnixNano())
	gw.maybeSweepTargetHealth()
	waitForCondition(t, func() bool { return atomic.LoadInt32(&hits) == 2 }, time.Second)
}

// 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{})
	release := 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()
		}
		<-release // never closes on its own; relies on the probe not reading
	}))
	// F6, feat/target-health review: previously relied on the PROBE'S OWN
	// 5s context timeout (targetTimeout below) to unblock this handler,
	// so srv.Close (deferred) silently added a real ~5s stall to every
	// run of this test. An explicit release channel, unblocked via
	// t.Cleanup BEFORE srv.Close (LIFO — see
	// TestMaybeSweepTargetHealth_ConcurrentTriggers_SingleFlight's own
	// comment for the identical ordering reasoning), makes this test
	// fast AND immune to a failing assertion ever leaving it blocked.
	t.Cleanup(srv.Close)
	t.Cleanup(func() { close(release) })

	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)
	}
}