Skip to content

Commit a0f2504

Browse files
theJCclaude
andauthored
Preserve query params in remote MCP server URLs (#4154)
* Preserve query params in remote MCP server URLs When a remote MCP server is registered with query parameters in the URL (e.g., ?toolsets=core,alerting,...), the transparent proxy was silently dropping them before forwarding requests to the upstream server. The remote server received only the base path and fell back to its default behavior, returning fewer capabilities than configured with no error or warning. Extract the raw query string alongside the path from the registration URL and merge it into every outbound request via a new WithRemoteRawQuery proxy option. Remote query params are prepended before any client-supplied params so operator-configured values take precedence (first-value-wins semantics). Raw string concatenation is used intentionally to preserve literal characters like commas that url.Values.Encode() would percent-encode. Also fix GenerateMCPServerURL to include query parameters in the URL returned to clients for SSE and streamable transports. Add a TransparentProxy.ListenerAddr() method to expose the bound address after start, enabling cross-package tests to use port 0. Add unit tests covering the proxy director merge logic, the HTTPTransport.Start() extraction path, and URL generation. Add an E2E test using a local mock upstream to cover the full CLI→proxy→upstream stack. Closes #4153 Signed-off-by: Jon Christiansen <467023+theJC@users.noreply.github.com> Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * Update per PR feedback --------- Signed-off-by: Jon Christiansen <467023+theJC@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
1 parent bc31b79 commit a0f2504

7 files changed

Lines changed: 443 additions & 23 deletions

File tree

pkg/transport/http.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,11 @@ func (t *HTTPTransport) Start(ctx context.Context) error {
252252
// paths so they reach the correct endpoint on the remote server.
253253
var remoteBasePath string
254254

255+
// remoteRawQuery holds the raw query string from the remote URL (e.g.,
256+
// "toolsets=core,alerting" from "https://mcp.example.com/mcp?toolsets=core,alerting").
257+
// This must be forwarded on every outbound request or it is silently dropped.
258+
var remoteRawQuery string
259+
255260
if t.remoteURL != "" {
256261
// For remote MCP servers, construct target URI from remote URL
257262
remoteURL, err := url.Parse(t.remoteURL)
@@ -267,9 +272,11 @@ func (t *HTTPTransport) Start(ctx context.Context) error {
267272
// The target URI only has scheme+host, so without this the remote path is lost.
268273
remoteBasePath = remoteURL.Path
269274

275+
remoteRawQuery = remoteURL.RawQuery
276+
270277
//nolint:gosec // G706: logging proxy port and remote URL from config
271278
slog.Debug("setting up transparent proxy to forward to remote URL",
272-
"port", t.proxyPort, "target", targetURI, "base_path", remoteBasePath)
279+
"port", t.proxyPort, "target", targetURI, "base_path", remoteBasePath, "raw_query", remoteRawQuery)
273280
} else {
274281
if t.containerName == "" {
275282
return transporterrors.ErrContainerNameNotSet
@@ -311,6 +318,7 @@ func (t *HTTPTransport) Start(ctx context.Context) error {
311318
if remoteBasePath != "" {
312319
proxyOptions = append(proxyOptions, transparent.WithRemoteBasePath(remoteBasePath))
313320
}
321+
proxyOptions = append(proxyOptions, transparent.WithRemoteRawQuery(remoteRawQuery))
314322

315323
// Create the transparent proxy
316324
t.proxy = transparent.NewTransparentProxyWithOptions(
Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package transport
5+
6+
import (
7+
"context"
8+
"fmt"
9+
"net/http"
10+
"net/http/httptest"
11+
"strings"
12+
"sync/atomic"
13+
"testing"
14+
"time"
15+
16+
"github.com/stretchr/testify/assert"
17+
"github.com/stretchr/testify/require"
18+
19+
"github.com/stacklok/toolhive/pkg/transport/proxy/transparent"
20+
"github.com/stacklok/toolhive/pkg/transport/types"
21+
)
22+
23+
// TestHTTPTransport_Start_RemoteURLQueryParams verifies that HTTPTransport.Start()
24+
// correctly extracts the raw query from the remoteURL and wires it into the
25+
// transparent proxy so every upstream request carries those query parameters.
26+
func TestHTTPTransport_Start_RemoteURLQueryParams(t *testing.T) {
27+
t.Parallel()
28+
29+
tests := []struct {
30+
name string
31+
remoteQuery string // query string appended to the remote registration URL
32+
expectedQuery string // raw query the upstream server should receive
33+
description string
34+
}{
35+
{
36+
name: "query params from registration URL are forwarded to upstream",
37+
remoteQuery: "toolsets=core,alerting",
38+
expectedQuery: "toolsets=core,alerting",
39+
description: "Datadog case: toolset selection params must reach the upstream server",
40+
},
41+
{
42+
name: "multiple query params are all forwarded to upstream",
43+
remoteQuery: "toolsets=core,alerting&version=2",
44+
expectedQuery: "toolsets=core,alerting&version=2",
45+
description: "Multiple params must all be forwarded, none dropped",
46+
},
47+
{
48+
name: "no query params — upstream receives empty query string",
49+
remoteQuery: "",
50+
expectedQuery: "",
51+
description: "Without configured query params, upstream receives an empty query string",
52+
},
53+
}
54+
55+
for _, tt := range tests {
56+
t.Run(tt.name, func(t *testing.T) {
57+
t.Parallel()
58+
59+
var receivedQuery atomic.Value
60+
61+
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
62+
receivedQuery.Store(r.URL.RawQuery)
63+
w.Header().Set("Content-Type", "application/json")
64+
w.WriteHeader(http.StatusOK)
65+
_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"1","result":{"protocolVersion":"2024-11-05"}}`))
66+
}))
67+
defer upstream.Close()
68+
69+
remoteURL := upstream.URL + "/mcp"
70+
if tt.remoteQuery != "" {
71+
remoteURL += "?" + tt.remoteQuery
72+
}
73+
74+
// Use port 0 so the OS assigns a free port.
75+
transport := NewHTTPTransport(
76+
types.TransportTypeStreamableHTTP,
77+
LocalhostIPv4,
78+
0, // proxyPort: OS-assigned
79+
0, // targetPort: unused for remote
80+
nil, // deployer: nil for remote
81+
false, // debug
82+
"", // targetHost: unused for remote
83+
nil, // authInfoHandler
84+
nil, // prometheusHandler
85+
nil, // prefixHandlers
86+
"", // endpointPrefix
87+
false, // trustProxyHeaders
88+
)
89+
transport.SetRemoteURL(remoteURL)
90+
91+
ctx, cancel := context.WithCancel(context.Background())
92+
defer cancel()
93+
94+
require.NoError(t, transport.Start(ctx))
95+
defer func() {
96+
assert.NoError(t, transport.Stop(context.Background()))
97+
}()
98+
99+
// Retrieve the actual listening address from the underlying proxy.
100+
tp, ok := transport.proxy.(*transparent.TransparentProxy)
101+
require.True(t, ok, "proxy should be a TransparentProxy")
102+
addr := tp.ListenerAddr()
103+
require.NotEmpty(t, addr, "proxy should be listening")
104+
105+
// POST to the clean proxy URL (no query params) so only the
106+
// proxy-configured remoteRawQuery is the source of upstream query params.
107+
proxyURL := fmt.Sprintf("http://%s/mcp", addr)
108+
body := `{"jsonrpc":"2.0","method":"initialize","id":"1","params":{}}`
109+
req, err := http.NewRequest(http.MethodPost, proxyURL, strings.NewReader(body))
110+
require.NoError(t, err)
111+
req.Header.Set("Content-Type", "application/json")
112+
113+
client := &http.Client{Timeout: 5 * time.Second}
114+
resp, err := client.Do(req)
115+
require.NoError(t, err)
116+
defer resp.Body.Close()
117+
118+
assert.Equal(t, http.StatusOK, resp.StatusCode)
119+
120+
actualQuery, _ := receivedQuery.Load().(string)
121+
assert.Equal(t, tt.expectedQuery, actualQuery,
122+
"%s: upstream server received wrong query string", tt.description)
123+
})
124+
}
125+
}

pkg/transport/proxy/transparent/remote_path_test.go

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,123 @@ import (
1818
"github.com/stretchr/testify/require"
1919
)
2020

21+
// TestRemoteQueryForwarding verifies that the transparent proxy correctly
22+
// forwards query parameters from the remote URL configuration to every
23+
// outbound request.
24+
//
25+
// Scenario: remoteURL is https://mcp.datadoghq.com/mcp?toolsets=core,alerting
26+
// Without this fix the query params are silently dropped and the remote
27+
// server receives /mcp with no toolsets, returning only default tools.
28+
func TestRemoteQueryForwarding(t *testing.T) {
29+
t.Parallel()
30+
31+
tests := []struct {
32+
name string
33+
remoteRawQuery string // Query from registration URL
34+
clientRawQuery string // Additional query from client request
35+
expectedRawQuery string // Query that should arrive at the remote server
36+
description string
37+
}{
38+
{
39+
name: "remote query only, no client query",
40+
remoteRawQuery: "toolsets=core,alerting",
41+
clientRawQuery: "",
42+
expectedRawQuery: "toolsets=core,alerting",
43+
description: "Datadog case: remote query params forwarded when client sends none",
44+
},
45+
{
46+
name: "remote query merged with client query",
47+
remoteRawQuery: "toolsets=core,alerting",
48+
clientRawQuery: "session=abc",
49+
expectedRawQuery: "toolsets=core,alerting&session=abc",
50+
description: "Remote params take precedence, client params appended",
51+
},
52+
{
53+
name: "no remote query, client query preserved",
54+
remoteRawQuery: "",
55+
clientRawQuery: "session=abc",
56+
expectedRawQuery: "session=abc",
57+
description: "Without remote query, client query passes through unchanged",
58+
},
59+
{
60+
name: "no remote query and no client query",
61+
remoteRawQuery: "",
62+
clientRawQuery: "",
63+
expectedRawQuery: "",
64+
description: "No query params in either direction",
65+
},
66+
}
67+
68+
for _, tt := range tests {
69+
t.Run(tt.name, func(t *testing.T) {
70+
t.Parallel()
71+
72+
var receivedQuery atomic.Value
73+
74+
remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
75+
receivedQuery.Store(r.URL.RawQuery)
76+
w.Header().Set("Content-Type", "application/json")
77+
w.WriteHeader(http.StatusOK)
78+
_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"1","result":{"protocolVersion":"2024-11-05"}}`))
79+
}))
80+
defer remoteServer.Close()
81+
82+
parsedRemote, err := url.Parse(remoteServer.URL)
83+
require.NoError(t, err)
84+
targetURI := (&url.URL{
85+
Scheme: parsedRemote.Scheme,
86+
Host: parsedRemote.Host,
87+
}).String()
88+
89+
var opts []Option
90+
if tt.remoteRawQuery != "" {
91+
opts = append(opts, WithRemoteRawQuery(tt.remoteRawQuery))
92+
}
93+
94+
proxy := NewTransparentProxyWithOptions(
95+
"127.0.0.1", 0, targetURI,
96+
nil, nil, nil,
97+
false, true, "streamable-http",
98+
nil, nil,
99+
"", false,
100+
nil, // middlewares
101+
opts...,
102+
)
103+
104+
ctx := context.Background()
105+
err = proxy.Start(ctx)
106+
require.NoError(t, err)
107+
defer func() {
108+
assert.NoError(t, proxy.Stop(context.Background()))
109+
}()
110+
111+
addr := proxy.ListenerAddr()
112+
require.NotEmpty(t, addr)
113+
114+
proxyURL := fmt.Sprintf("http://%s/mcp", addr)
115+
if tt.clientRawQuery != "" {
116+
proxyURL += "?" + tt.clientRawQuery
117+
}
118+
119+
body := `{"jsonrpc":"2.0","method":"initialize","id":"1","params":{}}`
120+
req, err := http.NewRequest(http.MethodPost, proxyURL, strings.NewReader(body))
121+
require.NoError(t, err)
122+
req.Header.Set("Content-Type", "application/json")
123+
124+
client := &http.Client{Timeout: 5 * time.Second}
125+
resp, err := client.Do(req)
126+
require.NoError(t, err)
127+
defer resp.Body.Close()
128+
129+
assert.Equal(t, http.StatusOK, resp.StatusCode)
130+
131+
actualQuery, _ := receivedQuery.Load().(string)
132+
assert.Equal(t, tt.expectedRawQuery, actualQuery,
133+
"%s: remote server received wrong query string", tt.description)
134+
})
135+
}
136+
}
137+
21138
// TestRemotePathForwarding verifies that the transparent proxy correctly
22139
// forwards requests to the remote server's full path, not just the host.
23140
//

pkg/transport/proxy/transparent/transparent_proxy.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,12 @@ type TransparentProxy struct {
108108
// URI only contains the scheme and host.
109109
remoteBasePath string
110110

111+
// remoteRawQuery holds the raw query string from the remote URL (e.g.,
112+
// "toolsets=core,alerting" from "https://mcp.example.com/mcp?toolsets=core,alerting").
113+
// When set, it is merged into every outbound request so query parameters
114+
// from the original registration URL are never silently dropped.
115+
remoteRawQuery string
116+
111117
// Deprecated: trustProxyHeaders indicates whether to trust X-Forwarded-* headers (moved to SSEResponseProcessor)
112118
trustProxyHeaders bool
113119

@@ -178,6 +184,18 @@ func WithRemoteBasePath(basePath string) Option {
178184
}
179185
}
180186

187+
// WithRemoteRawQuery sets the raw query string from the remote URL.
188+
// When set, these query parameters are merged into every outbound request,
189+
// ensuring query parameters from the original registration URL are always forwarded.
190+
// Ignores empty strings; default (no query forwarding) will be used.
191+
func WithRemoteRawQuery(rawQuery string) Option {
192+
return func(p *TransparentProxy) {
193+
if rawQuery != "" {
194+
p.remoteRawQuery = rawQuery
195+
}
196+
}
197+
}
198+
181199
// withHealthCheckPingTimeout sets the health check ping timeout.
182200
// This is primarily useful for testing with shorter timeouts.
183201
// Ignores non-positive timeouts; default will be used.
@@ -505,6 +523,21 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
505523
pr.Out.URL.RawPath = ""
506524
}
507525

526+
// Merge query parameters from the remote URL into the outbound request.
527+
// Remote params are prepended so they appear first; most HTTP servers
528+
// adopt first-value-wins semantics for duplicate keys, ensuring operator
529+
// configured values (e.g., toolsets=core,alerting) take precedence over
530+
// any client-supplied params with the same key.
531+
// Raw string concatenation is intentional: url.Values.Encode() would
532+
// percent-encode characters like commas that some APIs expect as literals.
533+
if p.remoteRawQuery != "" {
534+
merged := p.remoteRawQuery
535+
if pr.Out.URL.RawQuery != "" {
536+
merged += "&" + pr.Out.URL.RawQuery
537+
}
538+
pr.Out.URL.RawQuery = merged
539+
}
540+
508541
// Inject OpenTelemetry trace propagation headers for downstream tracing
509542
if pr.Out.Context() != nil {
510543
otel.GetTextMapPropagator().Inject(pr.Out.Context(), propagation.HeaderCarrier(pr.Out.Header))
@@ -605,6 +638,15 @@ func (p *TransparentProxy) Start(ctx context.Context) error {
605638
return nil
606639
}
607640

641+
// ListenerAddr returns the network address the proxy is listening on.
642+
// Returns an empty string if the proxy has not been started.
643+
func (p *TransparentProxy) ListenerAddr() string {
644+
if p.listener == nil {
645+
return ""
646+
}
647+
return p.listener.Addr().String()
648+
}
649+
608650
// CloseListener closes the listener for the transparent proxy.
609651
func (p *TransparentProxy) CloseListener() error {
610652
if p.listener != nil {

0 commit comments

Comments
 (0)