traefik-llmgateway/mcp_a2a_test.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)
	}
}

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

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