Files
traefik-llmgateway/mcp_a2a_test.go
T
lukaszraczylo 4a70daa8b8 fix(mcp): stop duplicate tool calls and keep full tool metadata
Federated tools/call re-initializes a session only when the backend
signals a missing session: HTTP 400 or 404, or a JSON-RPC error that
names a session. A failing tool is no longer called twice.

tools/list now:
- passes each tool object through and rewrites only its name
- follows nextCursor up to a page cap and keeps the earlier pages if a
  later page fails
- skips tools with no name
- resolves name collisions to the server that tools/call routes to

The handshake sends notifications/initialized. tools/call forwards
_meta. Multi-line SSE data is joined. A2A target URLs are joined with
url.JoinPath and keep their query string. agents[].card must start with
"/". Logged target URLs and errors are scrubbed of credentials.

Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
2026-09-23 09:44:32 +01:00

1309 lines
49 KiB
Go

package traefikllmgateway
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"net/url"
"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, "/")
}
}
// --- F7 (review round 3, 2026-09): URL joining, not string concatenation ---
// TestBuildUpstreamTargetURL covers buildUpstreamTargetURL directly:
// good (ordinary join), bad (a target whose own path already ends in
// "/", the doubled-slash bug), and edge (a target carrying its own query
// string, which must survive alongside the request's own).
func TestBuildUpstreamTargetURL(t *testing.T) {
cases := []struct {
name string
targetURL string
rest string
requestRawQuery string
want string
}{
{
name: "ordinary join",
targetURL: "http://x.internal/mcp",
rest: "foo",
want: "http://x.internal/mcp/foo",
},
{
name: "target path already ends in / no longer doubles the slash",
targetURL: "http://x.internal/mcp/",
rest: "foo",
want: "http://x.internal/mcp/foo",
},
{
name: "empty rest proxies to the target root unchanged",
targetURL: "http://x.internal/mcp",
rest: "",
want: "http://x.internal/mcp",
},
{
name: "target's own query survives alongside the request's own",
targetURL: "http://x.internal/mcp?token=abc",
rest: "foo",
requestRawQuery: "q=1",
want: "http://x.internal/mcp/foo?q=1&token=abc",
},
{
name: "target's own query survives with no request query at all",
targetURL: "http://x.internal/mcp?token=abc",
rest: "foo",
want: "http://x.internal/mcp/foo?token=abc",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got, err := buildUpstreamTargetURL(tc.targetURL, tc.rest, tc.requestRawQuery)
if err != nil {
t.Fatalf("buildUpstreamTargetURL: %v", err)
}
if got != tc.want {
t.Errorf("buildUpstreamTargetURL(%q, %q, %q) = %q, want %q", tc.targetURL, tc.rest, tc.requestRawQuery, got, tc.want)
}
})
}
}
// TestHandleTargetProxy_TargetQueryString_MergedWithRequestQuery is the
// F7 end-to-end regression test: a target URL configured with its own
// query string (the only way TargetConfig can carry a credential today)
// must reach the upstream intact, merged with — not overwritten or
// corrupted by — the caller's own request query.
func TestHandleTargetProxy_TargetQueryString_MergedWithRequestQuery(t *testing.T) {
var gotURL string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotURL = r.URL.String()
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 + "/mcp?token=secret"}}
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)
}
req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools?q=1", 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())
}
u, err := url.Parse(gotURL)
if err != nil {
t.Fatalf("parse upstream-observed URL %q: %v", gotURL, err)
}
if u.Path != "/mcp/tools" {
t.Errorf("upstream path = %q, want %q", u.Path, "/mcp/tools")
}
got := u.Query()
if got.Get("token") != "secret" {
t.Errorf("upstream query token = %q, want %q (the target's own configured query must survive)", got.Get("token"), "secret")
}
if got.Get("q") != "1" {
t.Errorf("upstream query q = %q, want %q (the request's own query must survive too)", got.Get("q"), "1")
}
}
// --- 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,
},
{
// F8, review round 3, 2026-09: handleAgents builds the card
// URL as "/a2a/" + name + cardPath with no separator of its
// own — a Card that does not start with "/" concatenates
// straight onto the agent name instead of starting a new
// path segment.
name: "agent card not starting with / is a constructor error",
agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: "agent.json"}},
wantErr: true,
},
{
name: "agent card starting with / is valid",
agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: "/custom-card.json"}},
wantErr: false,
},
{
name: "agent card empty (default applies) is valid",
agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: ""}},
wantErr: false,
},
}
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"))
}
}