Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion internal/server/mcp_routing.go
Original file line number Diff line number Diff line change
Expand Up @@ -795,10 +795,19 @@ func (p *MCPProxyServer) RefreshPrompts() {
// scanner before they are ever registered (parity with tool-description
// poisoning detection).
upstreamPrompts = p.scanAggregatedPrompts(upstreamPrompts)
// Spec 100: rug-pull baseline. Detect pending/changed metadata vs the
// approved baseline and WITHHOLD those prompts from registration (compose
// in series after the TPA scan — scan detects poison, baseline detects
// change). A withheld prompt is absent from prompts/list and fails
// prompts/get natively; there is no runtime get-time gate.
approval := p.checkPromptApprovals(upstreamPrompts)
upstreamPrompts = filterBlockedPrompts(upstreamPrompts, approval.blocked)
all = buildAggregatedServerPrompts(builtins, upstreamPrompts, p.getPromptAggregated, p.logger)
p.logger.Info("refreshed prompts",
zap.Int("upstream_prompt_count", len(upstreamPrompts)),
zap.Int("total_prompt_count", len(all)))
zap.Int("total_prompt_count", len(all)),
zap.Int("withheld_pending", approval.pending),
zap.Int("withheld_changed", approval.changed))
} else {
// nil upstreamPrompts: the aggregation loop never runs, so the nil
// getPrompt is never invoked.
Expand Down
46 changes: 46 additions & 0 deletions internal/server/mcp_routing_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -948,6 +948,8 @@ func TestRefreshPrompts_AggregatesBuiltinsAndUpstream(t *testing.T) {
proxy, _ := createTestProxyWithRuntime(t, nil)
proxy.config.EnablePrompts = true
proxy.config.AggregateUpstreamPrompts = true // opt in to upstream aggregation
qOff := false
proxy.config.QuarantineEnabled = &qOff // spec 100: these tests verify aggregation, not the rug-pull baseline

upstreamSrv := newTestRefreshPromptsUpstream(t)
testServer := mcpserver.NewTestStreamableHTTPServer(upstreamSrv)
Expand Down Expand Up @@ -975,6 +977,46 @@ func TestRefreshPrompts_AggregatesBuiltinsAndUpstream(t *testing.T) {
require.Contains(t, prompts, "server-a__greeting", "aggregated upstream prompt must be registered under its direct name")
}

// TestRefreshPrompts_RugPullBaseline_WithholdsFirstSeen (spec 100) proves the
// full RefreshPrompts path withholds a first-seen prompt on a quarantine-
// enforced server, and that approving it registers it on the next refresh.
func TestRefreshPrompts_RugPullBaseline_WithholdsFirstSeen(t *testing.T) {
t.Setenv("MCPPROXY_DISABLE_OAUTH", "true")

proxy, _ := createTestProxyWithRuntime(t, []*config.ServerConfig{
// Server present in config as manual trust + quarantine on (the default).
{Name: "server-a", Protocol: "streamable-http", Enabled: true, TrustMode: string(config.TrustModeManual)},
})
proxy.config.EnablePrompts = true
proxy.config.AggregateUpstreamPrompts = true

upstreamSrv := newTestRefreshPromptsUpstream(t)
testServer := mcpserver.NewTestStreamableHTTPServer(upstreamSrv)
t.Cleanup(testServer.Close)

um := upstream.NewManager(zap.NewNop(), proxy.config, nil, secret.NewResolver(), nil)
t.Cleanup(func() { um.DisconnectAll() })
require.NoError(t, um.AddServerConfig("srv-a", &config.ServerConfig{
Name: "server-a", Protocol: "streamable-http", URL: testServer.URL, Enabled: true,
}))
client, ok := um.GetClient("srv-a")
require.True(t, ok)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
require.NoError(t, client.Connect(ctx))
proxy.upstreamManager = um

proxy.RefreshPrompts()
prompts := proxy.server.ListPrompts()
require.Contains(t, prompts, "setup-new-mcp-server", "built-ins are always registered")
require.NotContains(t, prompts, "server-a__greeting", "a first-seen prompt on a manual server is withheld (rug-pull baseline)")

// Approve it → it registers on the triggered refresh.
require.NoError(t, proxy.ApprovePrompt("server-a", "greeting", "tester"))
prompts = proxy.server.ListPrompts()
require.Contains(t, prompts, "server-a__greeting", "an approved prompt is registered")
}

// TestRefreshPrompts_AggregationDisabled_BuiltinsOnly verifies the default
// safe posture (PR #973 review): with EnablePrompts on but the opt-in
// aggregate_upstream_prompts flag off, RefreshPrompts serves ONLY the built-ins
Expand Down Expand Up @@ -1028,6 +1070,8 @@ func TestRefreshPrompts_ReadsLiveAggregateFlag(t *testing.T) {
live := rt.Config()
live.EnablePrompts = true
live.AggregateUpstreamPrompts = true
qOff := false
live.QuarantineEnabled = &qOff // spec 100: verify aggregation, not the rug-pull baseline

// Construction-time snapshot DISAGREES (aggregation off). If RefreshPrompts
// read p.config it would skip aggregation — the assertion below would fail.
Expand Down Expand Up @@ -1085,6 +1129,8 @@ func TestRefreshPrompts_PopulatesRoutingModeServers(t *testing.T) {
proxy, _ := createTestProxyWithRuntime(t, nil)
proxy.config.EnablePrompts = true
proxy.config.AggregateUpstreamPrompts = true // opt in to upstream aggregation
qOff := false
proxy.config.QuarantineEnabled = &qOff // spec 100: these tests verify aggregation, not the rug-pull baseline

upstreamSrv := newTestRefreshPromptsUpstream(t)
testServer := mcpserver.NewTestStreamableHTTPServer(upstreamSrv)
Expand Down
Loading
Loading