From 3690598475d06532abb62f59a7af7420953c14cf Mon Sep 17 00:00:00 2001 From: Roman Chernyak Date: Sat, 27 Jun 2026 22:58:54 +0200 Subject: [PATCH] feat(discovery): surface quarantined tools in retrieve_tools (locked, name-only) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Quarantined tools — both server-level quarantine and tool-level pending/changed approvals — are intentionally absent from the BM25 search index so their untrusted descriptions cannot expose a Tool Poisoning Attack (TPA) payload to the agent. As a side effect, retrieve_tools could only answer "no such tool" for a capability a quarantined server provides, even though the correct remediation is "ask the user to approve it". Add an opt-in second pass to retrieve_tools(include_disabled=true) that enumerates quarantined tools from authoritative state and returns NAME-ONLY locked entries (description and schema withheld) with a status + remediation: - server-level quarantine -> new status `server_quarantined` - tool-level pending/changed -> existing `pending_approval` The pass is fully self-contained — it does NOT alter the shared callability or classification helpers, so visible-tool counts and other consumers are unchanged. Specifically it: - matches the query against the tool NAME only (quarantined tools aren't in the index, so they can't be BM25-ranked); short keywords (>=2 chars, e.g. "ui"/"qa") are retained; - applies the same agent-scope + profile filtering as the callable path, via a shared `serverDiscoverable` helper (de-duplicates what were three inline copies down to one for the discovery surface); - dedups against tools the index loop already handled (a `seen` set), so the brief post-quarantine reindex window can't list a tool as both callable and locked, nor double-count it in the zero-result nudge; - skips tools also denied by operator config (enabled_tools/disabled_tools), which approval cannot unlock — so the agent isn't sent down a dead-end remediation; - prepends its matches so the shared min(limit,10) cap can't truncate them in favor of index hits. Quarantined tools remain non-callable: the call path already blocks server-quarantine and pending/changed before execution (handleQuarantinedToolCall), so no change to isToolCallable is needed — discovery only makes them VISIBLE, never callable. Tests: pending tool surfaced by name with withheld description + remediation; server-level branch (tool-names source injected); query-scoped exclusion; config-denied skipped; dedup against seen; short-keyword tokens; zero-result nudge; quarantined tool still blocked at the call path; remediation present. Co-Authored-By: Claude Opus 4.8 (1M context) --- internal/contracts/types.go | 19 +- internal/server/mcp.go | 183 ++++++++++++++++- .../server/mcp_quarantine_discovery_test.go | 191 ++++++++++++++++++ 3 files changed, 377 insertions(+), 16 deletions(-) create mode 100644 internal/server/mcp_quarantine_discovery_test.go diff --git a/internal/contracts/types.go b/internal/contracts/types.go index c8d706dee..779969b63 100644 --- a/internal/contracts/types.go +++ b/internal/contracts/types.go @@ -223,16 +223,21 @@ type Tool struct { } // DisabledToolStatus is the single machine-branchable reason a tool exists but -// is not callable (Spec 049). Exactly one value per locked tool, assigned by -// fixed first-match precedence (server-off → config → user → pending → unknown). +// is not callable (Spec 049). Exactly one value per locked tool. The +// classifier (classifyServerToolStatus) assigns the index-discoverable reasons +// by fixed first-match precedence (server-off → config → user → pending → +// unknown). DisabledStatusServerQuarantined is assigned separately by the +// quarantined-tool discovery pass (quarantined tools are never in the index), +// not by the classifier. type DisabledToolStatus = string const ( - DisabledStatusServerDisabled DisabledToolStatus = "server_disabled" - DisabledStatusByConfig DisabledToolStatus = "disabled_by_config" - DisabledStatusByUser DisabledToolStatus = "disabled_by_user" - DisabledStatusPendingApproval DisabledToolStatus = "pending_approval" - DisabledStatusUnknown DisabledToolStatus = "disabled_unknown" + DisabledStatusServerDisabled DisabledToolStatus = "server_disabled" + DisabledStatusServerQuarantined DisabledToolStatus = "server_quarantined" + DisabledStatusByConfig DisabledToolStatus = "disabled_by_config" + DisabledStatusByUser DisabledToolStatus = "disabled_by_user" + DisabledStatusPendingApproval DisabledToolStatus = "pending_approval" + DisabledStatusUnknown DisabledToolStatus = "disabled_unknown" ) // LockedToolEntry is the lean discovery shape for a non-callable tool returned diff --git a/internal/server/mcp.go b/internal/server/mcp.go index 14a3986c4..d8a48a33a 100644 --- a/internal/server/mcp.go +++ b/internal/server/mcp.go @@ -11,6 +11,7 @@ import ( "sync" "sync/atomic" "time" + "unicode" "github.com/smart-mcp-proxy/mcpproxy-go/internal/auth" "github.com/smart-mcp-proxy/mcpproxy-go/internal/cache" @@ -1160,6 +1161,17 @@ func (p *MCPProxyServer) handleRetrieveToolsWithMode(ctx context.Context, reques args["include_disabled"] = true } + // serverDiscoverable applies the agent-scope (Spec 049 FR-007) and profile + // (Spec 057) filters BEFORE classification so an agent never learns a tool + // exists on a server it cannot access. Shared by the callable-result loop + // and the quarantined-tool discovery pass below so the two never drift. + serverDiscoverable := func(serverName string) bool { + if enforceAgentScope && authCtx != nil && !authCtx.CanAccessServer(serverName) { + return false + } + return profileScope.Allows(serverName) + } + var callableResults []*config.SearchResult var disabledEntries []contracts.LockedToolEntry droppedCount := 0 @@ -1174,15 +1186,7 @@ func (p *MCPProxyServer) handleRetrieveToolsWithMode(ctx context.Context, reques } } - // Agent-scope filtering happens BEFORE classification (Spec 049 FR-007) - // so an agent never learns a locked tool exists on a server it cannot - // access. - if enforceAgentScope && !authCtx.CanAccessServer(serverName) { - continue - } - // Spec 057: profile filter — runs independently of agent-scope so that - // unauthenticated /mcp/p/ connections (AdminContext) are still filtered. - if !profileScope.Allows(serverName) { + if !serverDiscoverable(serverName) { continue } @@ -1203,6 +1207,34 @@ func (p *MCPProxyServer) handleRetrieveToolsWithMode(ctx context.Context, reques } results = callableResults + // Quarantined tools are deliberately absent from the search index: their + // untrusted descriptions/schemas are withheld so a Tool Poisoning Attack + // payload is never exposed to the agent. The loop above therefore can never + // surface them, and an agent searching for a capability a quarantined server + // provides would be told "no such tool". Enumerate those tools from + // authoritative state (server-level quarantine + tool-level pending/changed + // approvals) and prepend name-only locked entries (no description/schema) so + // the agent learns the capability EXISTS but needs the user's approval. + // + // `seen` dedupes against tools the loop already handled — primarily for the + // brief window after a runtime quarantine toggle when a tool may still + // linger in the index — so a tool is never both a callable result and a + // locked entry, and droppedCount isn't double-counted. Matches are PREPENDED + // so the shared min(limit,10) cap below can't truncate them away in favor of + // index hits. The count feeds the zero-result nudge even without the flag. + seen := make(map[string]bool, len(callableResults)+len(disabledEntries)) + for _, r := range callableResults { + seen[r.Tool.Name] = true + } + for _, e := range disabledEntries { + seen[e.Name] = true + } + quarantinedMatches := p.collectQuarantinedToolMatches(query, serverDiscoverable, seen, p.serverToolNames) + droppedCount += len(quarantinedMatches) + if includeDisabled { + disabledEntries = append(quarantinedMatches, disabledEntries...) + } + // Spec 035 F4: Resolve annotations for each result and apply annotation-based filtering // before building the MCP tool response. This allows agents to self-restrict discovery. annotationFilterActive := readOnlyOnly || excludeDestructive || excludeOpenWorld @@ -5205,6 +5237,9 @@ func disabledToolRemediation(status contracts.DisabledToolStatus) string { switch status { case contracts.DisabledStatusServerDisabled: return "Its server is disabled. Ask the user to enable the server first." + case contracts.DisabledStatusServerQuarantined: + return "Its server is quarantined for security review. Its tools cannot be called " + + "until the user reviews and approves the server in the mcpproxy UI or system tray." case contracts.DisabledStatusByConfig: return "Locked by operator policy in mcp_config.json (enabled_tools/disabled_tools). " + "The user cannot enable this from the UI; ask the operator to change the server config." @@ -5218,6 +5253,136 @@ func disabledToolRemediation(status contracts.DisabledToolStatus) string { } } +// maxQuarantinedMatches bounds how many quarantined tools the discovery +// second-pass collects before stopping. The response itself is capped lower +// (min(limit,10)); this just bounds work on the opt-in path. +const maxQuarantinedMatches = 50 + +// collectQuarantinedToolMatches finds tools that exist but are quarantined and +// whose name matches the query, returning lean locked entries (no description +// or schema — those are withheld because a quarantined tool's description is a +// potential Tool Poisoning Attack payload). It covers both quarantine layers: +// - server-level quarantine: every tool on a Quarantined server is locked +// (status server_quarantined); +// - tool-level quarantine: pending/changed approval records on a trusted +// server (status pending_approval). +// +// Quarantined tools are not in the search index, so this is the only path that +// can surface them to discovery. allowServer mirrors the agent-scope/profile +// filtering applied to the index results so an agent never learns a tool +// exists on a server it cannot access. +// +// Inputs: +// - allowServer / seen: the discovery scope filter and the set of "server:tool" +// keys the index loop already handled (callable or locked), so a tool is +// never surfaced twice. +// - toolNamesFor: resolves a server's live tool names (injected for testing; +// production passes p.serverToolNames). +// +// A tool that is also denied by operator config (enabled_tools/disabled_tools) +// is skipped rather than advertised as merely "pending approval": approving it +// in the UI would not make it callable, so surfacing it would send the agent +// down a remediation that can never succeed. +func (p *MCPProxyServer) collectQuarantinedToolMatches(query string, allowServer func(string) bool, seen map[string]bool, toolNamesFor func(string) []string) []contracts.LockedToolEntry { + tokens := queryTokens(query) + if len(tokens) == 0 { + return nil + } + + servers, err := p.storage.ListUpstreams() + if err != nil { + p.logger.Debug("collectQuarantinedToolMatches: failed to list upstreams", zap.Error(err)) + return nil + } + + var matches []contracts.LockedToolEntry + // add records a quarantined tool; returns false once the work cap is hit. + add := func(sc *config.ServerConfig, toolName string, status contracts.DisabledToolStatus) bool { + key := sc.Name + ":" + toolName + if seen[key] || !toolNameMatchesQuery(toolName, tokens) { + return true + } + if p.isToolConfigDenied(sc.Name, toolName, sc) { + return true // operator policy, not quarantine — don't advertise as approvable + } + seen[key] = true + matches = append(matches, contracts.LockedToolEntry{Name: key, Server: sc.Name, Status: status}) + return len(matches) < maxQuarantinedMatches + } + + for _, sc := range servers { + if sc == nil || !sc.Enabled || !allowServer(sc.Name) { + continue + } + + if sc.Quarantined { + for _, toolName := range toolNamesFor(sc.Name) { + if !add(sc, toolName, contracts.DisabledStatusServerQuarantined) { + return matches + } + } + continue + } + + approvals, aerr := p.storage.ListToolApprovals(sc.Name) + if aerr != nil { + continue + } + for _, a := range approvals { + if a == nil || a.Disabled { + continue + } + if a.Status != storage.ToolApprovalStatusPending && a.Status != storage.ToolApprovalStatusChanged { + continue + } + if !add(sc, a.ToolName, contracts.DisabledStatusPendingApproval) { + return matches + } + } + } + return matches +} + +// queryTokens lowercases the query and splits it into alphanumeric tokens, used +// for a lightweight name match against quarantined tools (which are not in the +// BM25 index and so can't be ranked normally). Tokens of length >= 2 are kept so +// short-but-meaningful capability keywords ("ci", "ui", "qa", "db", "os") still +// match; if the query has only single-character fields they are kept as a +// fallback rather than returning nothing. +func queryTokens(query string) []string { + fields := strings.FieldsFunc(strings.ToLower(query), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsDigit(r) + }) + tokens := make([]string, 0, len(fields)) + for _, f := range fields { + if len(f) >= 2 { + tokens = append(tokens, f) + } + } + if len(tokens) == 0 { + for _, f := range fields { + if f != "" { + tokens = append(tokens, f) + } + } + } + return tokens +} + +// toolNameMatchesQuery reports whether a quarantined tool's name is relevant to +// the query: any query token appears in the lowercased tool name. Tool names +// usually carry the capability keywords (e.g. "emulator_build_web"), so this is +// enough to surface "this exists but is quarantined" without the description. +func toolNameMatchesQuery(toolName string, tokens []string) bool { + lower := strings.ToLower(toolName) + for _, tok := range tokens { + if strings.Contains(lower, tok) { + return true + } + } + return false +} + func (p *MCPProxyServer) lookupToolAnnotations(serverName, toolName string) *config.ToolAnnotations { if p.mainServer == nil || p.mainServer.runtime == nil { return nil diff --git a/internal/server/mcp_quarantine_discovery_test.go b/internal/server/mcp_quarantine_discovery_test.go new file mode 100644 index 000000000..cd6f927de --- /dev/null +++ b/internal/server/mcp_quarantine_discovery_test.go @@ -0,0 +1,191 @@ +package server + +import ( + "context" + "encoding/json" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" + "github.com/smart-mcp-proxy/mcpproxy-go/internal/contracts" + "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" +) + +// Quarantined tools are intentionally absent from the search index (their +// untrusted descriptions are withheld to avoid Tool Poisoning Attack exposure), +// so retrieve_tools' index loop can never surface them. These tests cover the +// discovery second pass that enumerates them from authoritative state, the +// description-withholding, the config-denied precedence, the dedup/scope/short- +// keyword handling, and the fact that surfaced quarantined tools stay +// non-callable at the call path. + +func findDisabled(entries []contracts.LockedToolEntry, name string) (contracts.LockedToolEntry, bool) { + for _, e := range entries { + if e.Name == name { + return e, true + } + } + return contracts.LockedToolEntry{}, false +} + +// allowAll is the scope filter used in unit tests (no agent scope / profile). +func allowAll(string) bool { return true } + +// A tool-level-quarantined (pending) tool on a trusted server is surfaced by +// include_disabled when its NAME matches the query — description withheld, +// pending_approval status + remediation. +func TestQuarantineDiscovery_PendingToolSurfacedByName(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{Name: "android", Enabled: true})) + require.NoError(t, proxy.storage.SaveToolApproval(&storage.ToolApprovalRecord{ + ServerName: "android", ToolName: "emulator_build_web", + Status: storage.ToolApprovalStatusPending, + })) + + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{ + "query": "build emulator web", "limit": float64(10), "include_disabled": true, + } + result, err := proxy.handleRetrieveTools(context.Background(), req) + require.NoError(t, err) + resp := decodeRetrieve(t, result) + + entry, ok := findDisabled(resp.Disabled, "android:emulator_build_web") + require.True(t, ok, "quarantined tool must appear in disabled[] (got %+v)", resp.Disabled) + assert.Equal(t, contracts.DisabledStatusPendingApproval, entry.Status) + assert.Empty(t, entry.Description, "quarantined tool description must be withheld (TPA payload)") + assert.Contains(t, resp.Remediation, contracts.DisabledStatusPendingApproval) + assert.Empty(t, resp.Tools, "quarantined tool must not be a callable result") +} + +// The second pass is query-scoped: a quarantined tool whose name does not match +// the query is NOT dumped into the results. +func TestQuarantineDiscovery_NameMustMatchQuery(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{Name: "android", Enabled: true})) + require.NoError(t, proxy.storage.SaveToolApproval(&storage.ToolApprovalRecord{ + ServerName: "android", ToolName: "emulator_build_web", + Status: storage.ToolApprovalStatusChanged, + })) + + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{ + "query": "frobnicate widget", "limit": float64(10), "include_disabled": true, + } + result, err := proxy.handleRetrieveTools(context.Background(), req) + require.NoError(t, err) + resp := decodeRetrieve(t, result) + + _, ok := findDisabled(resp.Disabled, "android:emulator_build_web") + assert.False(t, ok, "non-matching quarantined tool must not be surfaced") +} + +// Even without include_disabled, a query that matches only a quarantined tool +// gets the one-line nudge (droppedCount is bumped by the second pass) without +// inlining any locked entries. +func TestQuarantineDiscovery_ZeroResultNudge(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{Name: "android", Enabled: true})) + require.NoError(t, proxy.storage.SaveToolApproval(&storage.ToolApprovalRecord{ + ServerName: "android", ToolName: "emulator_build_web", + Status: storage.ToolApprovalStatusPending, + })) + + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"query": "build emulator web"} + result, err := proxy.handleRetrieveTools(context.Background(), req) + require.NoError(t, err) + resp := decodeRetrieve(t, result) + + text := result.Content[0].(mcp.TextContent).Text + assert.Contains(t, text, "include_disabled:true") + assert.Contains(t, text, "locked") + assert.Nil(t, resp.Disabled, "nudge must not inline locked entries") +} + +// Server-level quarantine: every tool on a Quarantined server is surfaced as a +// name-only locked entry with server_quarantined status. serverToolNames needs +// a wired runtime, so the tool-names source is injected here to exercise the +// branch (the production call passes p.serverToolNames). +func TestQuarantineDiscovery_ServerLevelBranch(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{ + Name: "collar-emu", Enabled: true, Quarantined: true, + })) + toolNames := func(string) []string { return []string{"emulator_build_web", "unrelated_tool"} } + + matches := proxy.collectQuarantinedToolMatches("build emulator", allowAll, map[string]bool{}, toolNames) + + require.Len(t, matches, 1, "only the name-matching tool surfaces (got %+v)", matches) + assert.Equal(t, "collar-emu:emulator_build_web", matches[0].Name) + assert.Equal(t, contracts.DisabledStatusServerQuarantined, matches[0].Status) + assert.Empty(t, matches[0].Description, "description must be withheld") +} + +// A pending tool that is ALSO denied by operator config must NOT be advertised +// as pending_approval — approving it in the UI cannot make it callable. +func TestQuarantineDiscovery_ConfigDeniedSkipped(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{ + Name: "android", Enabled: true, DisabledTools: []string{"emulator_build_web"}, + })) + require.NoError(t, proxy.storage.SaveToolApproval(&storage.ToolApprovalRecord{ + ServerName: "android", ToolName: "emulator_build_web", + Status: storage.ToolApprovalStatusPending, + })) + + matches := proxy.collectQuarantinedToolMatches("build emulator web", allowAll, map[string]bool{}, proxy.serverToolNames) + assert.Empty(t, matches, "config-denied tool must not be surfaced as pending_approval") +} + +// Tools already handled by the index loop (passed in via `seen`) are not +// surfaced again by the second pass. +func TestQuarantineDiscovery_DedupAgainstSeen(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{Name: "android", Enabled: true})) + require.NoError(t, proxy.storage.SaveToolApproval(&storage.ToolApprovalRecord{ + ServerName: "android", ToolName: "emulator_build_web", + Status: storage.ToolApprovalStatusPending, + })) + + seen := map[string]bool{"android:emulator_build_web": true} + matches := proxy.collectQuarantinedToolMatches("build emulator", allowAll, seen, proxy.serverToolNames) + assert.Empty(t, matches, "already-seen tool must not be surfaced twice") +} + +// A quarantined tool stays non-callable at the call path even though discovery +// surfaces it. +func TestQuarantinedTool_VisibleButNotCallable(t *testing.T) { + proxy := createTestMCPProxyServer(t) + require.NoError(t, proxy.storage.SaveUpstreamServer(&config.ServerConfig{ + Name: "collar-emu", Enabled: true, Quarantined: true, + })) + + result := proxy.directToolCallabilityBlock(context.Background(), "collar-emu", "emulator_build_web", map[string]interface{}{}) + require.NotNil(t, result, "a quarantined server's tool call must be blocked") + var response map[string]interface{} + require.NoError(t, json.Unmarshal([]byte(result.Content[0].(mcp.TextContent).Text), &response)) + assert.Equal(t, "QUARANTINED_SERVER_BLOCKED", response["status"]) +} + +// Short capability keywords ("ui", "qa", ...) are retained so they can match +// quarantined tool names; an all-single-char query falls back rather than +// returning nothing. +func TestQueryTokens_ShortKeywords(t *testing.T) { + assert.Equal(t, []string{"ui", "tap"}, queryTokens("ui tap")) + assert.Equal(t, []string{"qa", "build"}, queryTokens("qa build!")) + assert.Equal(t, []string{"a"}, queryTokens("a"), "single-char fallback") + assert.Empty(t, queryTokens(" ")) +} + +// Every status the discovery path can emit must have a non-empty, actionable +// remediation string. +func TestDisabledToolRemediation_ServerQuarantined(t *testing.T) { + msg := disabledToolRemediation(contracts.DisabledStatusServerQuarantined) + assert.NotEmpty(t, msg) + assert.Contains(t, msg, "quarantined") + assert.Contains(t, msg, "approve") +}