fix(go-adk): forward allowedHeaders from incoming request to MCP tool calls (#1733)

Closes #1679

The allowedHeaders field on McpServer was being populated by the
translator and stored in HttpMcpServerConfig/SseMcpServerConfig, but the
Go runtime never used it. Headers from the incoming A2A request were
silently dropped even when the operator explicitly listed them in
allowedHeaders. The Python runtime handled this correctly already.

What changed:

- allowedRequestHeaders helper added to registry.go. It reads the
a2asrv.CallContext already present in the Go context and returns only
the header values whose names appear in the configured allowedHeaders
list (case-insensitive). No intermediate context key, no extra copy.
- headerRoundTripper gets an allowedHeaders field. On each outbound MCP
HTTP request the transport calls allowedRequestHeaders and sets the
matching headers. Static headers from the server spec are applied last
so they always take precedence, mirroring the Python runtime behaviour.
- CreateToolsets now passes AllowedHeaders from each MCP server config
into the transport params.
- executor.go no longer needs to touch headers at all. The A2A server
already stores the CallContext in the Go context before Execute is
called, so the round tripper picks it up automatically.

The change is fully backwards compatible. If allowedHeaders is empty, no
extra headers are forwarded.

Tests use a2asrv.NewRequestMeta and a2asrv.WithCallContext to build the
test context, matching what the real A2A server does. Coverage includes:
normal forwarding, static-header-wins override, no A2A context,
non-allowed header filter, and empty allowed list.

Signed-off-by: mesutoezdil <mesudozdil@gmail.com>

---------

Signed-off-by: mesutoezdil <mesudozdil@gmail.com>
Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io>
Co-authored-by: Eitan Yarmush <eitan.yarmush@solo.io>

Author: mesutoezdil

go/adk/pkg/mcp/registry.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/a2aproject/a2a-go/a2asrv"
 	"github.com/go-logr/logr"
 	"github.com/kagent-dev/kagent/go/api/adk"
 	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
@@ -21,12 +22,47 @@ const (
 	defaultTimeout = 30 * time.Minute
 )
 
+// allowedRequestHeaders reads the incoming A2A request metadata from ctx and
+// returns only the header key/value pairs whose names appear in allowed.
+// It reads directly from the A2A CallContext that is already present in the Go
+// context, avoiding a redundant copy.
+//
+// Lookup relies on RequestMeta.Get which already does a case-insensitive O(1)
+// lookup (NewRequestMeta lowercases keys at construction). Keys in the result
+// preserve the casing from the allowed list so the MCP server sees the header
+// names the operator configured. When a header has multiple values only the
+// first one is forwarded; additional values are intentionally dropped.
+func allowedRequestHeaders(ctx context.Context, allowed []string) map[string]string {
+	if len(allowed) == 0 {
+		return nil
+	}
+	callCtx, ok := a2asrv.CallContextFrom(ctx)
+	if !ok {
+		return nil
+	}
+	meta := callCtx.RequestMeta()
+	if meta == nil {
+		return nil
+	}
+	result := make(map[string]string)
+	for _, name := range allowed {
+		if vals, ok := meta.Get(name); ok && len(vals) > 0 && vals[0] != "" {
+			result[name] = vals[0]
+		}
+	}
+	if len(result) == 0 {
+		return nil
+	}
+	return result
+}
+
 // mcpServerParams groups connection parameters for an MCP server,
 // reducing parameter sprawl across createTransport / initializeToolSet.
 type mcpServerParams struct {
 	URL                   string
 	Headers               map[string]string
-	ServerType            string // "http" or "sse"
+	AllowedHeaders        []string // header names to forward from incoming request
+	ServerType            string   // "http" or "sse"
 	Timeout               *float64
 	SseReadTimeout        *float64
 	TLSInsecureSkipVerify *bool
@@ -46,6 +82,7 @@ func CreateToolsets(ctx context.Context, httpTools []adk.HttpMcpServerConfig, ss
 		params := mcpServerParams{
 			URL:                   httpTool.Params.Url,
 			Headers:               httpTool.Params.Headers,
+			AllowedHeaders:        httpTool.AllowedHeaders,
 			ServerType:            "http",
 			Timeout:               httpTool.Params.Timeout,
 			SseReadTimeout:        httpTool.Params.SseReadTimeout,
@@ -65,6 +102,7 @@ func CreateToolsets(ctx context.Context, httpTools []adk.HttpMcpServerConfig, ss
 		params := mcpServerParams{
 			URL:                   sseTool.Params.Url,
 			Headers:               sseTool.Params.Headers,
+			AllowedHeaders:        sseTool.AllowedHeaders,
 			ServerType:            "sse",
 			Timeout:               sseTool.Params.Timeout,
 			SseReadTimeout:        sseTool.Params.SseReadTimeout,
@@ -162,10 +200,11 @@ func createTransport(ctx context.Context, params mcpServerParams) (mcpsdk.Transp
 	}
 
 	var httpTransport http.RoundTripper = baseTransport
-	if len(params.Headers) > 0 {
+	if len(params.Headers) > 0 || len(params.AllowedHeaders) > 0 {
 		httpTransport = &headerRoundTripper{
-			base:    baseTransport,
-			headers: params.Headers,
+			base:           baseTransport,
+			headers:        params.Headers,
+			allowedHeaders: params.AllowedHeaders,
 		}
 	}
 
@@ -190,17 +229,35 @@ func createTransport(ctx context.Context, params mcpServerParams) (mcpsdk.Transp
 	return mcpTransport, nil
 }
 
-// headerRoundTripper wraps an http.RoundTripper to add custom headers to all requests.
+// headerRoundTripper wraps an http.RoundTripper to add custom headers to all
+// requests. It supports two sources of headers:
+//   - headers: static key/value pairs configured on the MCP server spec
+//   - allowedHeaders: header names to forward from the incoming A2A request;
+//     values are read on each call via allowedRequestHeaders directly from the
+//     A2A CallContext that is already present in the Go context.
+//
+// Static headers take precedence: if an allowed header has the same name as a
+// static header, the static value wins.
 type headerRoundTripper struct {
-	base    http.RoundTripper
-	headers map[string]string
+	base           http.RoundTripper
+	headers        map[string]string
+	allowedHeaders []string // header names (case-insensitive) to forward from A2A context
 }
 
 func (rt *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	req = req.Clone(req.Context())
+
+	// Forward allowed headers from the incoming A2A request first so that
+	// static headers can override them if there is a name collision.
+	for k, v := range allowedRequestHeaders(req.Context(), rt.allowedHeaders) {
+		req.Header.Set(k, v)
+	}
+
+	// Apply static headers (override any dynamic ones with the same name).
 	for key, value := range rt.headers {
 		req.Header.Set(key, value)
 	}
+
 	return rt.base.RoundTrip(req)
 }
 

go/adk/pkg/mcp/registry_test.go

@@ -0,0 +1,254 @@
+package mcp
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/a2aproject/a2a-go/a2asrv"
+)
+
+// a2aCtx builds a context that carries an A2A CallContext with the given headers.
+// Keys are stored case-insensitively by NewRequestMeta, matching the behaviour
+// of a real A2A server.
+func a2aCtx(headers map[string][]string) context.Context {
+	meta := a2asrv.NewRequestMeta(headers)
+	ctx, _ := a2asrv.WithCallContext(context.Background(), meta)
+	return ctx
+}
+
+// TestAllowedRequestHeaders_ForwardsMatchingHeaders verifies that headers listed
+// in allowedHeaders are forwarded when they are present in the A2A CallContext.
+func TestAllowedRequestHeaders_ForwardsMatchingHeaders(t *testing.T) {
+	t.Parallel()
+	var capturedAuth, capturedCustom, capturedStatic string
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedAuth = r.Header.Get("Authorization")
+		capturedCustom = r.Header.Get("X-Custom")
+		capturedStatic = r.Header.Get("X-Static")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	ctx := a2aCtx(map[string][]string{
+		"Authorization": {"Bearer token123"},
+		"X-Custom":      {"custom-value"},
+		"X-Ignored":     {"should-not-appear"},
+	})
+
+	rt := &headerRoundTripper{
+		base:           http.DefaultTransport,
+		headers:        map[string]string{"X-Static": "static-value"},
+		allowedHeaders: []string{"Authorization", "X-Custom"},
+	}
+
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if capturedAuth != "Bearer token123" {
+		t.Errorf("Authorization: got %q, want %q", capturedAuth, "Bearer token123")
+	}
+	if capturedCustom != "custom-value" {
+		t.Errorf("X-Custom: got %q, want %q", capturedCustom, "custom-value")
+	}
+	if capturedStatic != "static-value" {
+		t.Errorf("X-Static: got %q, want %q", capturedStatic, "static-value")
+	}
+}
+
+// TestAllowedRequestHeaders_StaticOverridesDynamic verifies that a statically
+// configured header wins over the same header forwarded from the A2A request.
+func TestAllowedRequestHeaders_StaticOverridesDynamic(t *testing.T) {
+	t.Parallel()
+	var capturedAuth string
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedAuth = r.Header.Get("Authorization")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	ctx := a2aCtx(map[string][]string{
+		"Authorization": {"Bearer incoming"},
+	})
+
+	rt := &headerRoundTripper{
+		base:           http.DefaultTransport,
+		headers:        map[string]string{"Authorization": "Bearer static"},
+		allowedHeaders: []string{"Authorization"},
+	}
+
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if capturedAuth != "Bearer static" {
+		t.Errorf("Authorization: got %q, want %q", capturedAuth, "Bearer static")
+	}
+}
+
+// TestAllowedRequestHeaders_NoA2AContext verifies that no headers are forwarded
+// when the context does not carry an A2A CallContext.
+func TestAllowedRequestHeaders_NoA2AContext(t *testing.T) {
+	t.Parallel()
+	var capturedAuth string
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedAuth = r.Header.Get("Authorization")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	rt := &headerRoundTripper{
+		base:           http.DefaultTransport,
+		allowedHeaders: []string{"Authorization"},
+	}
+
+	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, srv.URL, nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if capturedAuth != "" {
+		t.Errorf("Authorization should be empty without A2A context, got %q", capturedAuth)
+	}
+}
+
+// TestAllowedRequestHeaders_IgnoresNonAllowed verifies that headers not listed
+// in allowedHeaders are not forwarded even if they appear in the A2A request.
+func TestAllowedRequestHeaders_IgnoresNonAllowed(t *testing.T) {
+	t.Parallel()
+	var capturedIgnored string
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedIgnored = r.Header.Get("X-Ignored")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	ctx := a2aCtx(map[string][]string{
+		"X-Ignored": {"should-not-appear"},
+	})
+
+	rt := &headerRoundTripper{
+		base:           http.DefaultTransport,
+		allowedHeaders: []string{"Authorization"},
+	}
+
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip failed: %v", err)
+	}
+	resp.Body.Close()
+
+	if capturedIgnored != "" {
+		t.Errorf("X-Ignored should not be forwarded, got %q", capturedIgnored)
+	}
+}
+
+// TestAllowedRequestHeaders_EmptyAllowedList verifies that allowedRequestHeaders
+// returns nil immediately when the allowed list is empty.
+func TestAllowedRequestHeaders_EmptyAllowedList(t *testing.T) {
+	t.Parallel()
+	ctx := a2aCtx(map[string][]string{
+		"Authorization": {"Bearer token"},
+	})
+
+	got := allowedRequestHeaders(ctx, nil)
+	if got != nil {
+		t.Errorf("expected nil for empty allowed list, got %v", got)
+	}
+
+	got = allowedRequestHeaders(ctx, []string{})
+	if got != nil {
+		t.Errorf("expected nil for empty allowed list, got %v", got)
+	}
+}
+
+// TestAllowedRequestHeaders_CaseInsensitiveLookup verifies that matching between
+// the configured allowedHeaders and the incoming request headers is case-insensitive
+// regardless of which side is lowercased or uppercased.
+func TestAllowedRequestHeaders_CaseInsensitiveLookup(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name     string
+		incoming map[string][]string
+		allowed  []string
+		wantKey  string
+		wantVal  string
+	}{
+		{
+			name:     "allowed lowercase, incoming capitalized",
+			incoming: map[string][]string{"Authorization": {"Bearer x"}},
+			allowed:  []string{"authorization"},
+			wantKey:  "authorization",
+			wantVal:  "Bearer x",
+		},
+		{
+			name:     "allowed capitalized, incoming lowercase",
+			incoming: map[string][]string{"authorization": {"Bearer y"}},
+			allowed:  []string{"Authorization"},
+			wantKey:  "Authorization",
+			wantVal:  "Bearer y",
+		},
+		{
+			name:     "mixed case both sides",
+			incoming: map[string][]string{"X-Trace-Id": {"abc"}},
+			allowed:  []string{"x-trace-id"},
+			wantKey:  "x-trace-id",
+			wantVal:  "abc",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := a2aCtx(tc.incoming)
+			got := allowedRequestHeaders(ctx, tc.allowed)
+			if got[tc.wantKey] != tc.wantVal {
+				t.Errorf("got[%q] = %q, want %q (full map: %v)", tc.wantKey, got[tc.wantKey], tc.wantVal, got)
+			}
+		})
+	}
+}
+
+// TestAllowedRequestHeaders_MultiValueFirstWins documents the behaviour for headers
+// that arrive with multiple values: only the first one is forwarded. If a use case
+// ever needs all values, the helper signature will have to change.
+func TestAllowedRequestHeaders_MultiValueFirstWins(t *testing.T) {
+	t.Parallel()
+	ctx := a2aCtx(map[string][]string{
+		"X-Forwarded-For": {"1.2.3.4", "5.6.7.8", "9.10.11.12"},
+	})
+	got := allowedRequestHeaders(ctx, []string{"X-Forwarded-For"})
+	if got["X-Forwarded-For"] != "1.2.3.4" {
+		t.Errorf("expected first value 1.2.3.4, got %q", got["X-Forwarded-For"])
+	}
+}
+
+// TestAllowedRequestHeaders_ReturnsNilWhenNoMatches verifies that the helper returns
+// nil rather than an empty map when the allowed list has entries but none of them
+// appear in the request metadata.
+func TestAllowedRequestHeaders_ReturnsNilWhenNoMatches(t *testing.T) {
+	t.Parallel()
+	ctx := a2aCtx(map[string][]string{
+		"X-Something-Else": {"value"},
+	})
+	got := allowedRequestHeaders(ctx, []string{"Authorization", "X-Trace-Id"})
+	if got != nil {
+		t.Errorf("expected nil when no allowed headers are present, got %v", got)
+	}
+}