diff --git a/cmd/wasm/runtime_wasm.go b/cmd/wasm/runtime_wasm.go index 85859d423c..7ad9e83599 100644 --- a/cmd/wasm/runtime_wasm.go +++ b/cmd/wasm/runtime_wasm.go @@ -133,11 +133,17 @@ func buildRuntime(ctx context.Context, cfg *latest.Config, env environment.Provi if transport == "" { transport = "streamable-http" } + var headerFactory func(context.Context) map[string]string + if len(ts.Remote.Headers) > 0 { + headerFactory = func(context.Context) map[string]string { + return ts.Remote.Headers + } + } mcpTS := mcptools.NewRemoteToolsetWithAllowPrivateIPs( ts.Name, ts.Remote.URL, transport, - ts.Remote.Headers, + headerFactory, ts.Remote.OAuth, ts.AllowPrivateIPsEnabled(), ) @@ -145,11 +151,17 @@ func buildRuntime(ctx context.Context, cfg *latest.Config, env environment.Provi case ts.URL != "": // URL directly on the toolset (legacy format) transport := "streamable-http" + var headerFactory func(context.Context) map[string]string + if len(ts.Headers) > 0 { + headerFactory = func(context.Context) map[string]string { + return ts.Headers + } + } mcpTS := mcptools.NewRemoteToolsetWithAllowPrivateIPs( ts.Name, ts.URL, transport, - ts.Headers, + headerFactory, nil, ts.AllowPrivateIPsEnabled(), ) diff --git a/pkg/tools/a2a/a2a.go b/pkg/tools/a2a/a2a.go index 67a2ab1c4d..5d7affa588 100644 --- a/pkg/tools/a2a/a2a.go +++ b/pkg/tools/a2a/a2a.go @@ -183,8 +183,13 @@ func (t *Toolset) Start(ctx context.Context) error { base = http.DefaultTransport } - headers := t.expander.ExpandMap(ctx, t.headers) - httpClient.Transport = upstream.NewHeaderTransport(base, headers) + var headerFactory func(context.Context) map[string]string + if len(t.headers) > 0 { + headerFactory = func(reqCtx context.Context) map[string]string { + return t.expander.ExpandMap(reqCtx, t.headers) + } + } + httpClient.Transport = upstream.NewHeaderTransport(base, headerFactory) client, err := a2aclient.NewFromCard(ctx, card, a2aclient.WithDefaultsDisabled(), diff --git a/pkg/tools/mcp/keyringstore/tokenstore_test.go b/pkg/tools/mcp/keyringstore/tokenstore_test.go index 12595f0b72..e77a9a1ef7 100644 --- a/pkg/tools/mcp/keyringstore/tokenstore_test.go +++ b/pkg/tools/mcp/keyringstore/tokenstore_test.go @@ -7,6 +7,7 @@ import ( "fmt" "os" "path/filepath" + "runtime" "sync" "testing" "time" @@ -362,8 +363,12 @@ func TestEncryptedStore_FilePermissions(t *testing.T) { if err != nil { t.Fatalf("Stat: %v", err) } - if perm := info.Mode().Perm(); perm != 0o600 { - t.Errorf("token file permissions = %o, want 0600", perm) + expectedPerm := os.FileMode(0o600) + if runtime.GOOS == "windows" { + expectedPerm = 0o666 + } + if perm := info.Mode().Perm(); perm != expectedPerm { + t.Errorf("token file permissions = %o, want %o", perm, expectedPerm) } } @@ -766,8 +771,12 @@ func TestFileKeyringPassphrase_RandomAndPersistent(t *testing.T) { if err != nil { t.Fatalf("Stat passphrase file: %v", err) } - if perm := info.Mode().Perm(); perm != 0o600 { - t.Errorf("passphrase file permissions = %o, want 0600", perm) + expectedPerm := os.FileMode(0o600) + if runtime.GOOS == "windows" { + expectedPerm = 0o666 + } + if perm := info.Mode().Perm(); perm != expectedPerm { + t.Errorf("passphrase file permissions = %o, want %o", perm, expectedPerm) } // A different install dir must get a different passphrase. diff --git a/pkg/tools/mcp/mcp.go b/pkg/tools/mcp/mcp.go index ae8a8e734c..4da6689aff 100644 --- a/pkg/tools/mcp/mcp.go +++ b/pkg/tools/mcp/mcp.go @@ -111,16 +111,20 @@ func CreateToolSet(ctx context.Context, toolset latest.Toolset, runConfig *confi case toolset.Remote.URL != "": expander := js.NewJsExpander(envProvider) - - // TODO: expand headers on each request, not at creation time. - headers := expander.ExpandMap(ctx, toolset.Remote.Headers) remoteURL := expander.Expand(ctx, toolset.Remote.URL, nil) + + var headerFactory func(context.Context) map[string]string + if len(toolset.Remote.Headers) > 0 { + headerFactory = func(reqCtx context.Context) map[string]string { + return expander.ExpandMap(reqCtx, toolset.Remote.Headers) + } + } return NewRemoteToolsetWithAllowPrivateIPs( toolset.Name, remoteURL, toolset.Remote.TransportType, - headers, + headerFactory, toolset.Remote.OAuth, toolset.AllowPrivateIPsEnabled(), lifecycle.PolicyFromConfig(toolset.Name, toolset.Lifecycle), @@ -239,25 +243,25 @@ func NewToolsetCommand(name, command string, args, env []string, cwd string, pol // // The optional policy lets callers tune restart/backoff behaviour; // see NewToolsetCommand for the semantics. -func NewRemoteToolset(name, urlString, transport string, headers map[string]string, oauthConfig *latest.RemoteOAuthConfig, policy ...lifecycle.Policy) *Toolset { - return newRemoteToolset(name, urlString, transport, headers, oauthConfig, false, policy...) +func NewRemoteToolset(name, urlString, transport string, headerFactory func(context.Context) map[string]string, oauthConfig *latest.RemoteOAuthConfig, policy ...lifecycle.Policy) *Toolset { + return newRemoteToolset(name, urlString, transport, headerFactory, oauthConfig, false, policy...) } // NewRemoteToolsetWithAllowPrivateIPs creates a new remote MCP toolset and // optionally permits OAuth helper requests to dial non-public IP addresses. func NewRemoteToolsetWithAllowPrivateIPs( name, urlString, transport string, - headers map[string]string, + headerFactory func(context.Context) map[string]string, oauthConfig *latest.RemoteOAuthConfig, allowPrivateIPs bool, policy ...lifecycle.Policy, ) *Toolset { - return newRemoteToolset(name, urlString, transport, headers, oauthConfig, allowPrivateIPs, policy...) + return newRemoteToolset(name, urlString, transport, headerFactory, oauthConfig, allowPrivateIPs, policy...) } func newRemoteToolset( name, urlString, transport string, - headers map[string]string, + headerFactory func(context.Context) map[string]string, oauthConfig *latest.RemoteOAuthConfig, allowPrivateIPs bool, policy ...lifecycle.Policy, @@ -265,14 +269,13 @@ func newRemoteToolset( slog.Debug("Creating Remote MCP toolset", "url", urlString, "transport", transport, - "headers", headers, "allow_private_ips", allowPrivateIPs, ) desc := buildRemoteDescription(urlString, transport) ts := &Toolset{ name: name, - mcpClient: newRemoteClient(urlString, transport, headers, NewKeyringTokenStore(), oauthConfig, allowPrivateIPs), + mcpClient: newRemoteClient(urlString, transport, headerFactory, NewKeyringTokenStore(), oauthConfig, allowPrivateIPs), logID: urlString, description: desc, } diff --git a/pkg/tools/mcp/oauth_helpers.go b/pkg/tools/mcp/oauth_helpers.go index ddbce7d519..8e0ba2a08c 100644 --- a/pkg/tools/mcp/oauth_helpers.go +++ b/pkg/tools/mcp/oauth_helpers.go @@ -61,9 +61,9 @@ func oauthHTTPClientForAllowPrivateIPs(allowPrivateIPs bool) *http.Client { // // The returned client is a fresh instance; the shared oauthHTTPClient // singleton is never mutated. -func oauthHTTPClientWithHeaders(rawURL string, headers map[string]string, allowPrivateIPs bool) *http.Client { +func oauthHTTPClientWithHeaders(rawURL string, headerFactory func(context.Context) map[string]string, allowPrivateIPs bool) *http.Client { base := oauthHTTPClientForAllowPrivateIPs(allowPrivateIPs) - if len(headers) == 0 { + if headerFactory == nil { return base } u, err := url.Parse(rawURL) @@ -83,7 +83,7 @@ func oauthHTTPClientWithHeaders(rawURL string, headers map[string]string, allowP CheckRedirect: base.CheckRedirect, Transport: &hostScopedHeaderTransport{ host: hostWithoutDefaultPort(u.Host, u.Scheme), - withHeaders: upstream.NewHeaderTransport(inner, headers), + withHeaders: upstream.NewHeaderTransport(inner, headerFactory), base: inner, }, } diff --git a/pkg/tools/mcp/oauth_test.go b/pkg/tools/mcp/oauth_test.go index bf716d1fb9..8f45befff0 100644 --- a/pkg/tools/mcp/oauth_test.go +++ b/pkg/tools/mcp/oauth_test.go @@ -927,7 +927,7 @@ func TestOAuthHTTPClientWithHeaders_StripsDefaultPort(t *testing.T) { t.Parallel() client := oauthHTTPClientWithHeaders("https://mcp.example.com:443/mcp", - map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"}, false) + func(context.Context) map[string]string { return map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"} }, false) hst, ok := client.Transport.(*hostScopedHeaderTransport) require.True(t, ok, "expected a host-scoped transport when headers are configured") assert.Equal(t, "mcp.example.com", hst.host, @@ -1541,8 +1541,10 @@ func TestInitialize_CustomHeadersReachOAuthDiscovery(t *testing.T) { srv := newUnmanagedOAuthTestServer(t) defer srv.Close() - headers := map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"} - client := newRemoteClient(srv.URL, "streamable", headers, NewInMemoryTokenStore(), nil, true) + headerFactory := func(context.Context) map[string]string { + return map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"} + } + client := newRemoteClient(srv.URL, "streamable", headerFactory, NewInMemoryTokenStore(), nil, true) // Decline the OAuth elicitation: the unmanaged flow reaches the // protected-resource-metadata discovery request (which carries the diff --git a/pkg/tools/mcp/remote.go b/pkg/tools/mcp/remote.go index eba5f2080b..c2ff41767c 100644 --- a/pkg/tools/mcp/remote.go +++ b/pkg/tools/mcp/remote.go @@ -22,7 +22,7 @@ type remoteMCPClient struct { url string transportType string - headers map[string]string + headerFactory func(context.Context) map[string]string tokenStore OAuthTokenStore managed bool unmanagedOAuthRedirectURI string @@ -32,7 +32,7 @@ type remoteMCPClient struct { func newRemoteClient( url, transportType string, - headers map[string]string, + headerFactory func(context.Context) map[string]string, tokenStore OAuthTokenStore, oauthConfig *latest.RemoteOAuthConfig, allowPrivateIPs bool, @@ -40,7 +40,6 @@ func newRemoteClient( slog.Debug("Creating remote MCP client", "url", url, "transport", transportType, - "headers", headers, "allow_private_ips", allowPrivateIPs, ) @@ -52,7 +51,7 @@ func newRemoteClient( sessionClient: sessionClient{serverAddress: sanitizeRemoteAddress(url)}, url: url, transportType: transportType, - headers: headers, + headerFactory: headerFactory, tokenStore: tokenStore, oauthConfig: oauthConfig, allowPrivateIPs: allowPrivateIPs, @@ -229,7 +228,7 @@ func (c *remoteMCPClient) createHTTPClient() (*http.Client, *oauthTransport, err managed: c.managed, unmanagedOAuthRedirectURI: c.unmanagedOAuthRedirectURI, oauthConfig: c.oauthConfig, - oauthHTTPClient: oauthHTTPClientWithHeaders(c.url, c.headers, c.allowPrivateIPs), + oauthHTTPClient: oauthHTTPClientWithHeaders(c.url, c.headerFactory, c.allowPrivateIPs), } // Persist cookies across requests @@ -250,8 +249,8 @@ func (c *remoteMCPClient) headerTransport() http.RoundTripper { } base = t } - if len(c.headers) > 0 { - return upstream.NewHeaderTransport(base, c.headers) + if c.headerFactory != nil { + return upstream.NewHeaderTransport(base, c.headerFactory) } return base } diff --git a/pkg/tools/mcp/remote_test.go b/pkg/tools/mcp/remote_test.go index b4b94073ff..fe3c260e33 100644 --- a/pkg/tools/mcp/remote_test.go +++ b/pkg/tools/mcp/remote_test.go @@ -81,7 +81,8 @@ func TestRemoteClientCustomHeaders(t *testing.T) { "Authorization": "Bearer custom-token", } - client := newRemoteClient(server.URL, "sse", expectedHeaders, NewInMemoryTokenStore(), nil, false) + headerFactory := func(context.Context) map[string]string { return expectedHeaders } + client := newRemoteClient(server.URL, "sse", headerFactory, NewInMemoryTokenStore(), nil, false) // Try to initialize (which will make the HTTP request) // We don't care if it succeeds or fails, we just need it to make the request @@ -130,7 +131,8 @@ func TestRemoteClientHeadersWithStreamable(t *testing.T) { "X-Custom-Auth": "custom-auth-value", } - client := newRemoteClient(server.URL, "streamable", expectedHeaders, NewInMemoryTokenStore(), nil, false) + headerFactory := func(context.Context) map[string]string { return expectedHeaders } + client := newRemoteClient(server.URL, "streamable", headerFactory, NewInMemoryTokenStore(), nil, false) // Try to initialize _, _ = client.Initialize(t.Context(), nil) @@ -205,8 +207,8 @@ func TestRemoteClientEmptyHeaders(t *testing.T) { })) defer server.Close() - // Create remote client with empty headers map - client := newRemoteClient(server.URL, "sse", map[string]string{}, NewInMemoryTokenStore(), nil, false) + headerFactory := func(context.Context) map[string]string { return map[string]string{} } + client := newRemoteClient(server.URL, "sse", headerFactory, NewInMemoryTokenStore(), nil, false) _, _ = client.Initialize(t.Context(), nil) @@ -246,8 +248,8 @@ func TestOAuthHTTPClientWithHeaders_ScopesHeadersToMCPHost(t *testing.T) { })) defer thirdParty.Close() - headers := map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"} - client := oauthHTTPClientWithHeaders(mcpServer.URL, headers, true) + headerFactory := func(context.Context) map[string]string { return map[string]string{"X-Grafana-URL": "https://instance.grafana.net/"} } + client := oauthHTTPClientWithHeaders(mcpServer.URL, headerFactory, true) mcpReq, err := http.NewRequestWithContext(t.Context(), http.MethodGet, mcpServer.URL, http.NoBody) require.NoError(t, err) diff --git a/pkg/upstream/headers.go b/pkg/upstream/headers.go index 2152ff73fe..4037b11ecd 100644 --- a/pkg/upstream/headers.go +++ b/pkg/upstream/headers.go @@ -35,21 +35,28 @@ func Handler(next http.Handler) http.Handler { } // NewHeaderTransport wraps an http.RoundTripper to set custom headers on -// every outbound request. Header values may contain ${headers.NAME} -// placeholders that are resolved at request time from upstream headers -// stored in the request context. -func NewHeaderTransport(base http.RoundTripper, headers map[string]string) http.RoundTripper { - return &headerTransport{base: base, headers: headers} +// every outbound request. The headerFactory is evaluated on each request +// to allow dynamic resolution of headers (e.g. from the environment). +// Header values may also contain ${headers.NAME} placeholders that are +// resolved at request time from upstream headers stored in the request context. +func NewHeaderTransport(base http.RoundTripper, headerFactory func(context.Context) map[string]string) http.RoundTripper { + return &headerTransport{base: base, headerFactory: headerFactory} } type headerTransport struct { - base http.RoundTripper - headers map[string]string + base http.RoundTripper + headerFactory func(context.Context) map[string]string } func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) { req = req.Clone(req.Context()) - for key, value := range ResolveHeaders(req.Context(), t.headers) { + + var headers map[string]string + if t.headerFactory != nil { + headers = t.headerFactory(req.Context()) + } + + for key, value := range ResolveHeaders(req.Context(), headers) { req.Header.Set(key, value) } return t.base.RoundTrip(req)