Files
traefik-llmgateway/mcp_a2a_test.go
T
lukaszraczylo f5d9014b13 test(target-health): deterministic yaegi-check and non-blocking sweep tests
exerciseTargetHealth's traffic-sourced assertion (mcpProbeServerName
must read healthy/traffic, F5 review) raced its own trigger: GET
/admin/api/targets was the first-ever sweep-capable route this
harness hit, so the background sweep it started could beat the
response encode and overwrite source from traffic to probe under
Yaegi's interpreter overhead. exerciseTargetHealthProbeWarmup now
runs first, before exerciseHandler, against a dedicated, never-
proxied MCP target — its health can only ever be written by the
probe, so polling it to healthy/probe both proves the interpreted
active-probe path end to end and spends this gateway's one
probeInterval-bounded sweep before any traffic runs, making every
later target-health assertion in this harness deterministic.
probeInterval moves from 10s to 60s so no later exercise can trigger
a second sweep mid-run. TARGETHEALTH-1's expected state also updates
to unknown, reflecting F4 (a2d152b1): mcpProbeServerName's only
observation is a failed tools/call, so it has never once succeeded.

mcp_a2a_test.go's ClientCanceled_RecordsNothing test replaces its
30ms-cancel-vs-200ms-sleep timing margin with a channel handshake
(cancel only once the handler is genuinely reached, then release
it), removing any wall-clock dependence.

Claude-Session: https://claude.ai/code/session_01EebV5jyWvYocByNQ3oss1w
2026-09-03 02:46:04 +01:00

1182 lines
45 KiB
Go

package traefikllmgateway
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
// --- unit tests: targetRoute ---
func TestTargetRoute(t *testing.T) {
tests := []struct {
path string
prefix string
wantName string
wantRest string
wantOK bool
}{
{"/mcp/myserver/tools/list", "mcp", "myserver", "tools/list", true},
{"/mcp/myserver", "mcp", "myserver", "", true},
{"/mcp/myserver/", "mcp", "myserver", "", true},
{"/mcp", "mcp", "", "", false},
{"/mcp/", "mcp", "", "", false},
{"/a2a/agent1/.well-known/agent-card.json", "a2a", "agent1", ".well-known/agent-card.json", true},
{"/a2a/agent1", "a2a", "agent1", "", true},
{"/notmcp/foo", "mcp", "", "", false},
{"/", "mcp", "", "", false},
{"", "mcp", "", "", false},
}
for _, tt := range tests {
t.Run(tt.prefix+tt.path, func(t *testing.T) {
gotName, gotRest, gotOK := targetRoute(tt.path, tt.prefix)
if gotOK != tt.wantOK || gotName != tt.wantName || gotRest != tt.wantRest {
t.Errorf("targetRoute(%q, %q) = (%q, %q, %v), want (%q, %q, %v)",
tt.path, tt.prefix, gotName, gotRest, gotOK, tt.wantName, tt.wantRest, tt.wantOK)
}
})
}
}
// --- registry filtering: GET /v1/mcp/servers, GET /v1/agents ---
func TestHandleMCPServers_GroupFiltered_SortedByName(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{
"charlie": {URL: "http://mcp-charlie.internal"},
"alpha": {URL: "http://mcp-alpha.internal"},
"bravo": {URL: "http://mcp-bravo.internal"},
}
cfg.Groups = map[string]*GroupConfig{"limited": {MCPServers: []string{"alpha", "charlie"}}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "limited", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/v1/mcp/servers", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
var got struct {
Servers []struct {
Name string `json:"name"`
URL string `json:"url"`
} `json:"servers"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
t.Fatalf("decode response: %v", err)
}
if len(got.Servers) != 2 {
t.Fatalf("servers = %+v, want 2 entries (alpha, charlie; bravo denied by group)", got.Servers)
}
if got.Servers[0].Name != "alpha" || got.Servers[0].URL != "/mcp/alpha" {
t.Errorf("servers[0] = %+v, want name=alpha url=/mcp/alpha", got.Servers[0])
}
if got.Servers[1].Name != "charlie" || got.Servers[1].URL != "/mcp/charlie" {
t.Errorf("servers[1] = %+v, want name=charlie url=/mcp/charlie", got.Servers[1])
}
}
func TestHandleAgents_GroupFiltered_SortedByName_CardPath(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Agents = map[string]*AgentConfig{
"zeta": {URL: "http://agent-zeta.internal"},
"delta": {URL: "http://agent-delta.internal", Card: "/custom/card.json"},
}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/v1/agents", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
var got struct {
Agents []struct {
Name string `json:"name"`
URL string `json:"url"`
Card string `json:"card"`
} `json:"agents"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
t.Fatalf("decode response: %v", err)
}
if len(got.Agents) != 2 {
t.Fatalf("agents = %+v, want 2 entries", got.Agents)
}
// sorted: delta before zeta
if got.Agents[0].Name != "delta" || got.Agents[0].Card != "/a2a/delta/custom/card.json" {
t.Errorf("agents[0] = %+v, want name=delta card=/a2a/delta/custom/card.json", got.Agents[0])
}
if got.Agents[1].Name != "zeta" || got.Agents[1].Card != "/a2a/zeta/.well-known/agent-card.json" {
t.Errorf("agents[1] = %+v, want name=zeta card=/a2a/zeta/.well-known/agent-card.json (default cardPath)", got.Agents[1])
}
}
func TestHandleMCPServers_Unauthenticated_Returns401(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
for _, path := range []string{"/v1/mcp/servers", "/v1/agents"} {
t.Run(path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, path, nil)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusUnauthorized {
t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String())
}
})
}
}
// --- target proxy: strips gateway auth, injects nothing ---
func TestHandleTargetProxy_StripsGatewayAuth_NoUpstreamInjection(t *testing.T) {
var gotAuth, gotAPIKeyHeader, gotConnection string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAuth = r.Header.Get("Authorization")
gotAPIKeyHeader = r.Header.Get("x-api-key")
gotConnection = r.Header.Get("Connection")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("mcp-ok"))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
req.Header.Set("Connection", "keep-alive")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != "mcp-ok" {
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
}
if gotAuth != "" {
t.Errorf("upstream saw Authorization %q, want stripped and nothing injected (in-cluster target)", gotAuth)
}
if gotAPIKeyHeader != "" {
t.Errorf("upstream saw x-api-key %q, want stripped", gotAPIKeyHeader)
}
if gotConnection != "" {
t.Errorf("upstream saw hop-by-hop Connection header %q, want stripped", gotConnection)
}
}
// TestHandleTargetProxy_ClientAPIKeyHeaderStripped_NoUpstreamInjection
// proves the client's gateway credential, presented via x-api-key rather
// than Authorization, never reaches an MCP target. handleTargetProxy
// injects no credential at all (unlike a provider adapter's injectAuth),
// so if gatewayCredentialHeaders' X-Api-Key entry were ever dropped from
// the strip set, the client's own key would pass straight through
// copyHeadersExcept to the upstream target (I1: the sibling test above
// authenticates via Authorization and never sends x-api-key at all, so it
// never exercised this strip entry).
func TestHandleTargetProxy_ClientAPIKeyHeaderStripped_NoUpstreamInjection(t *testing.T) {
var gotAPIKeyHeader string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAPIKeyHeader = r.Header.Get("x-api-key")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("mcp-ok"))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "gateway-key-must-not-leak"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil)
// Gateway auth presented via x-api-key only (no Authorization header) —
// presentedKey falls back to x-api-key when Authorization is absent.
req.Header.Set("x-api-key", "gateway-key-must-not-leak")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != "mcp-ok" {
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
}
if gotAPIKeyHeader != "" {
t.Errorf("upstream saw x-api-key = %q, want empty (MCP target proxy injects nothing; the client's gateway credential must be stripped)", gotAPIKeyHeader)
}
}
// TestHandleTargetProxy_DangerousHeadersStripped is the MCP/A2A-target
// half of the security+performance audit's (2026-08-22) additive header
// strip — proxyUpstream applies dangerousClientHeaders/
// dangerousClientHeaderPrefixes on BOTH proxy paths, so an MCP target
// must never see a client-spoofed identity/forwarding header, exactly
// like a native passthrough provider.
func TestHandleTargetProxy_DangerousHeadersStripped(t *testing.T) {
var got http.Header
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
got = r.Header.Clone()
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
req.Header.Set("Forwarded", "for=203.0.113.1")
req.Header.Set("X-Forwarded-For", "203.0.113.1")
req.Header.Set("X-Auth-Request-Email", "[email protected]")
req.Header.Set("X-Remote-User", "spoofed-admin")
req.Header.Set("X-Remote-Groups", "admin")
req.Header.Set("Cookie", "session=stolen")
req.Header.Set("X-Client-Custom", "keep-me")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
for _, hdr := range []string{"Forwarded", "X-Forwarded-For", "X-Auth-Request-Email", "X-Remote-User", "X-Remote-Groups", "Cookie"} {
if v := got.Get(hdr); v != "" {
t.Errorf("upstream MCP target saw %s = %q, want stripped", hdr, v)
}
}
if got.Get("X-Client-Custom") != "keep-me" {
t.Error("upstream did not see ordinary header X-Client-Custom — strip must not be a full allowlist inversion")
}
}
// --- SSE-safe streaming through the target proxy ---
func TestHandleTargetProxy_SSE_FlushesIncrementally(t *testing.T) {
continueCh := 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)
fl, ok := w.(http.Flusher)
if !ok {
t.Fatal("fake upstream ResponseWriter does not support Flush")
}
_, _ = w.Write([]byte("data: chunk-1\n\n"))
fl.Flush()
<-continueCh
_, _ = w.Write([]byte("data: chunk-2\n\n"))
fl.Flush()
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/stream", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
sw := newSyncFlushWriter()
done := make(chan struct{})
go func() {
h.ServeHTTP(sw, req)
close(done)
}()
select {
case <-sw.writeCh:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for the first chunk to reach the client")
}
_, writesAfterFirst, _, bodyAfterFirst := sw.snapshot()
if writesAfterFirst != 1 {
t.Fatalf("writes after first chunk = %d, want exactly 1 (chunks must not coalesce)", writesAfterFirst)
}
if !strings.Contains(bodyAfterFirst, "chunk-1") || strings.Contains(bodyAfterFirst, "chunk-2") {
t.Fatalf("body after first chunk = %q, want chunk-1 only", bodyAfterFirst)
}
close(continueCh)
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for ServeHTTP to finish")
}
status, writes, flushes, body := sw.snapshot()
if status != http.StatusOK {
t.Fatalf("status = %d, want 200", status)
}
if writes < 2 {
t.Fatalf("want >=2 writes (one per upstream chunk), got %d", writes)
}
if flushes < writes {
t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes)
}
if !strings.Contains(body, "chunk-1") || !strings.Contains(body, "chunk-2") {
t.Errorf("body missing streamed chunks, got %q", body)
}
}
// --- authz and validation checks ---
func TestHandleTargetProxy_UnknownName_Returns404(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}}
cfg.Agents = map[string]*AgentConfig{"agent1": {URL: "http://agent1.internal"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
for _, path := range []string{"/mcp/nosuchserver/tools/list", "/a2a/nosuchagent/tasks/send"} {
t.Run(path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, path, nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusNotFound {
t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String())
}
})
}
}
func TestHandleTargetProxy_Unauthenticated_Returns401(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
// Even a name that does not exist must be rejected as 401 before 404 —
// auth runs before the unknown-target check.
req := httptest.NewRequest(http.MethodGet, "/mcp/nosuchserver/tools/list", nil)
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusUnauthorized {
t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String())
}
}
func TestHandleTargetProxy_GroupDenies_Returns403(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{
"alpha": {URL: "http://mcp-alpha.internal"},
"beta": {URL: "http://mcp-beta.internal"},
}
cfg.Groups = map[string]*GroupConfig{"limited": {MCPServers: []string{"alpha"}}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "limited", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/beta/tools/list", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusForbidden {
t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String())
}
}
func TestHandleTargetProxy_TraversalPath_Returns400(t *testing.T) {
upstreamCalled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalled = true
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
for _, path := range []string{
"/mcp/alpha/../secret",
"/mcp/alpha/..%2f..%2fsecret",
} {
t.Run(path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, path, nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String())
}
})
}
if upstreamCalled {
t.Error("upstream must never be called for a rejected traversal path")
}
}
func TestHandleTargetProxy_UpgradeHeader_Returns501(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/realtime", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusNotImplemented {
t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String())
}
}
// --- empty rest proxies to target root ---
func TestHandleTargetProxy_EmptyRest_ProxiesToRoot(t *testing.T) {
var gotPath string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if gotPath != "/" {
t.Errorf("upstream path = %q, want %q (empty rest proxies to target root)", gotPath, "/")
}
}
// --- dead upstream maps to 502, same as native passthrough ---
func TestHandleTargetProxy_DeadUpstream_Returns502(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
deadURL := srv.URL
srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: deadURL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusBadGateway {
t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String())
}
}
// --- a2a proxies distinctly from mcp ---
func TestHandleTargetProxy_A2A_ProxiesToAgent(t *testing.T) {
var gotPath string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotPath = r.URL.Path
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("a2a-ok"))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Agents = map[string]*AgentConfig{"agent1": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := httptest.NewRequest(http.MethodPost, "/a2a/agent1/tasks/send", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != "a2a-ok" {
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
}
if gotPath != "/tasks/send" {
t.Errorf("upstream path = %q, want /tasks/send", gotPath)
}
}
// --- rate limiting: rejected-before-upstream requests still count, and a
// breached limit is enforced before the upstream is ever called ---
// TestHandleTargetProxy_RateLimited_Returns429WithRetryAfter proves the
// target proxy enforces the caller's group requestsPerMinute limit: the
// first request within the window succeeds, the second is refused with
// 429 and a Retry-After header, and the upstream is only ever called
// once.
func TestHandleTargetProxy_RateLimited_Returns429WithRetryAfter(t *testing.T) {
upstreamCalls := 0
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls++
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"limited": {Limits: &LimitsConfig{RequestsPerMinute: 1}}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "limited", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
req := func() *http.Request {
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/x", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
return req
}
rec1 := httptest.NewRecorder()
h.ServeHTTP(rec1, req())
if rec1.Code != http.StatusOK {
t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
}
rec2 := httptest.NewRecorder()
h.ServeHTTP(rec2, req())
if rec2.Code != http.StatusTooManyRequests {
t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String())
}
if rec2.Header().Get("Retry-After") == "" {
t.Error("want a Retry-After header on the 429 response")
}
if upstreamCalls != 1 {
t.Errorf("upstream calls = %d, want exactly 1 (the rejected second request must never reach the upstream)", upstreamCalls)
}
}
// --- accounting: a target request counts only against request-rate
// counters, never token/cost, even for a JSON response shaped like a
// usage-bearing one ---
// TestHandleTargetProxy_JSONResponse_OnlyRequestCounterMoves proves the
// target proxy never tees or parses a response for usage accounting
// (unlike native passthrough): a JSON body carrying usage-shaped fields
// leaves the caller's token/cost counters untouched, while the
// request-rate counter checkAndCount always increments still moves.
func TestHandleTargetProxy_JSONResponse_OnlyRequestCounterMoves(t *testing.T) {
const respBody = `{"model":"gpt-native-x","usage":{"prompt_tokens":7,"completion_tokens":3}}`
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(respBody))
}))
defer srv.Close()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}},
}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
if rec.Body.String() != respBody {
t.Errorf("body = %q, want verbatim upstream body", rec.Body.String())
}
tokIn, ok := gw.limiter.getCounter("user", "alice", metricTokIn, windowDay, time.Now())
if !ok || tokIn != 0 {
t.Errorf("user tokin/day counter = %d (ok=%v), want 0 — target proxy must never account response usage", tokIn, ok)
}
tokOut, ok := gw.limiter.getCounter("user", "alice", metricTokOut, windowDay, time.Now())
if !ok || tokOut != 0 {
t.Errorf("user tokout/day counter = %d (ok=%v), want 0 — target proxy must never account response usage", tokOut, ok)
}
cost, ok := gw.limiter.getCounter("user", "alice", metricCost, windowDay, time.Now())
if !ok || cost != 0 {
t.Errorf("user cost/day counter = %d (ok=%v), want 0 — target proxy must never account response cost", cost, ok)
}
reqCount, ok := gw.limiter.getCounter("user", "alice", metricReq, windowMin, time.Now())
if !ok || reqCount != 1 {
t.Errorf("user request/min counter = %d (ok=%v), want 1 — the request itself is still accounted", reqCount, ok)
}
// withTotalScope (handleTargetProxy, mcp_a2a.go) must have appended
// the synthetic total scope too (v0.2 data-layer task): counted like
// any other scope, still never token/cost-accounted here.
totalReq, ok := gw.limiter.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, time.Now())
if !ok || totalReq != 1 {
t.Errorf("total request/min counter = %d (ok=%v), want 1", totalReq, ok)
}
// Feature A (v0.22): handleTargetProxy/proxyUpstream never wrap the
// request context with an attemptRecorder — an MCP/A2A target proxy
// attempt against the "openai" MCP server's own upstream must never be
// mistaken for provider traffic, even though a provider of the same
// name happens to be configured too.
provAttempts, _ := gw.limiter.getCounter(kindProvider, "openai", metricProvAttempt, windowDay, time.Now())
if provAttempts != 0 {
t.Errorf("provider attempts/day = %d, want 0 — the MCP target proxy must never record provider accounting", provAttempts)
}
}
// --- config validation at construction ---
func TestNewGateway_TargetURLValidation(t *testing.T) {
tests := []struct {
mcpServers map[string]*TargetConfig
agents map[string]*AgentConfig
name string
wantErr bool
}{
{
name: "valid mcp server url",
mcpServers: map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal:8080"}},
wantErr: false,
},
{
name: "valid https agent url",
agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal"}},
wantErr: false,
},
{
name: "malformed mcp url (control character)",
mcpServers: map[string]*TargetConfig{"alpha": {URL: "http://exa\x7fmple.com"}},
wantErr: true,
},
{
name: "disallowed scheme for mcp server",
mcpServers: map[string]*TargetConfig{"alpha": {URL: "ftp://mcp-alpha.internal"}},
wantErr: true,
},
{
name: "disallowed scheme for agent",
agents: map[string]*AgentConfig{"agent1": {URL: "ws://agent1.internal"}},
wantErr: true,
},
{
name: "malformed agent url (invalid port)",
agents: map[string]*AgentConfig{"agent1": {URL: "http://agent1.internal:notaport"}},
wantErr: true,
},
{
name: "nil mcpServers config value is a constructor error, not a panic",
mcpServers: map[string]*TargetConfig{"alpha": nil},
wantErr: true,
},
{
name: "nil agents config value is a constructor error, not a panic",
agents: map[string]*AgentConfig{"agent1": nil},
wantErr: true,
},
{
name: "mcpServers name reserved (v1) is a constructor error",
mcpServers: map[string]*TargetConfig{"v1": {URL: "http://mcp-alpha.internal:8080"}},
wantErr: true,
},
{
name: "mcpServers name reserved (mcp) is a constructor error",
mcpServers: map[string]*TargetConfig{"mcp": {URL: "http://mcp-alpha.internal:8080"}},
wantErr: true,
},
{
name: "agents name reserved (a2a) is a constructor error",
agents: map[string]*AgentConfig{"a2a": {URL: "https://agent1.internal"}},
wantErr: true,
},
{
name: "mcpServers name with invalid character is a constructor error",
mcpServers: map[string]*TargetConfig{"has a space": {URL: "http://mcp-alpha.internal:8080"}},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = tt.mcpServers
cfg.Agents = tt.agents
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if tt.wantErr {
if err == nil {
t.Fatal("want error, got nil")
}
if h != nil {
t.Fatal("want nil handler on error")
}
return
}
if err != nil {
t.Fatalf("New: %v", err)
}
})
}
}
// --- Feature B (v0.21): per-target request accounting ---
func TestTargetScopeKind(t *testing.T) {
cases := []struct {
routingKind string
want string
}{
{targetKindMCP, "mcp"},
{targetKindAgent, scopeKindAgent},
}
for _, tc := range cases {
t.Run(tc.routingKind, func(t *testing.T) {
if got := targetScopeKind(tc.routingKind); got != tc.want {
t.Errorf("targetScopeKind(%q) = %q, want %q", tc.routingKind, got, tc.want)
}
})
}
}
// TestHandleTargetProxy_PerTargetCounters_AttributedAndCollisionSafe proves
// handleTargetProxy attributes a proxied request to name's own per-target
// scope (Feature B, v0.21) at every window countTargetRequest tracks
// (min/hour/day/month), for both an MCP server and an A2A agent target,
// AND that this never collides with a same-named user counter: both cases
// deliberately name the target "alice" — identical to the calling user's
// own name — so a bug that dropped kind from the counter key (windowKey,
// limits.go) would show up as the user's own counters silently absorbing
// the target's, or vice versa.
func TestHandleTargetProxy_PerTargetCounters_AttributedAndCollisionSafe(t *testing.T) {
cases := []struct {
wantScopeKind string
routingKind string
targetName string
path string
}{
{wantScopeKind: targetKindMCP, routingKind: targetKindMCP, targetName: "alice", path: "/mcp/alice/tools/list"},
{wantScopeKind: scopeKindAgent, routingKind: targetKindAgent, targetName: "alice", path: "/a2a/alice/tasks/send"},
}
for _, tc := range cases {
t.Run(tc.routingKind, func(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"}}
if tc.routingKind == targetKindMCP {
cfg.MCPServers = map[string]*TargetConfig{tc.targetName: {URL: srv.URL}}
} else {
cfg.Agents = map[string]*AgentConfig{tc.targetName: {URL: srv.URL}}
}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
req := httptest.NewRequest(http.MethodGet, tc.path, nil)
req.Header.Set("Authorization", "Bearer sk-alice")
rec := httptest.NewRecorder()
h.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
now := time.Now()
for _, window := range []string{windowMin, windowHour, windowDay, windowMonth} {
got, gotOK := gw.limiter.getCounter(tc.wantScopeKind, tc.targetName, metricReq, window, now)
if !gotOK || got != 1 {
t.Errorf("%s/%s req:%s counter = %d (ok=%v), want 1", tc.wantScopeKind, tc.targetName, window, got, gotOK)
}
}
// Collision guard: the caller's own USER scope (kind="user")
// must stay independent of the TARGET scope above despite
// sharing the identical id string "alice" — proving kind, not
// id alone, is what makes a counter key unique.
userReqMin, minOK := gw.limiter.getCounter("user", "alice", metricReq, windowMin, now)
if !minOK || userReqMin != 1 {
t.Errorf("user/alice req:min counter = %d (ok=%v), want 1 (its own scope, unaffected by the target scope of the same id)", userReqMin, minOK)
}
userReqMonth, monthOK := gw.limiter.getCounter("user", "alice", metricReq, windowMonth, now)
if !monthOK || userReqMonth != 0 {
t.Errorf("user/alice req:month counter = %d (ok=%v), want 0 (checkAndCount never writes a month window for user/group scopes; nonzero here would mean the target scope's month write leaked into it)", userReqMonth, monthOK)
}
})
}
}
// --- 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. F6, feat/target-health review: uses
// a channel handshake (cancel only once the handler has genuinely been
// reached, then release it) instead of a 30ms-cancel-vs-200ms-sleep
// wall-clock margin, which left no slack under load and could either
// cancel before the request reached the handler or fail to prove the
// cancel raced a still-in-flight request.
func TestHandleTargetProxy_TargetHealth_ClientCanceled_RecordsNothing(t *testing.T) {
reached := make(chan struct{})
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
close(reached)
<-release
w.WriteHeader(http.StatusOK)
}))
// Registered in this order so t.Cleanup's LIFO order unblocks the
// handler (closing release) BEFORE srv.Close runs — see
// TestMaybeSweepTargetHealth_ConcurrentTriggers_SingleFlight's own
// comment (target_health_test.go) for the identical reasoning.
t.Cleanup(srv.Close)
t.Cleanup(func() { close(release) })
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())
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil).WithContext(ctx)
req.Header.Set("Authorization", "Bearer sk-alice")
done := make(chan struct{})
go func() {
h.ServeHTTP(httptest.NewRecorder(), req)
close(done)
}()
<-reached
cancel()
<-done
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"))
}
}