fix(mcp): eliminate spurious grant-revoked errors (tool filter + session reset + system-user bypass) (#87)

* feat(mcp): filter tools at registration + detect FastMCP session reset with force-reconnect

Two foundational MCP reliability improvements:

(1) Tool allow/deny filtering at BridgeTool registration: Previously the runtime grant-check at execute time surfaced "grant revoked" errors when the LLM called a tool it wasn't allowed to call. Filter upfront in both the per-agent registration path (manager_connect.go) and per-user registration path (loop_mcp_user.go). Adds tool_filter.go with IsToolAllowed + tests. This eliminates the "registered then runtime-denied" loop.

(2) FastMCP session reset detection + force-reconnect: FastMCP/Python mcp servers reject tools/call as "invalid during session initialization" when the upstream session lifecycle resets but our pool still holds the old Mcp-Session-Id. Detector matches three known phrasings (FastMCP, mcp-go, mcp-go transport ErrSessionTerminated) in session_reset.go. On detection, BridgeTool.Execute requests a force-reconnect via atomic CAS dedup so N concurrent failing calls collapse to one reconnect. Health loops skip ping while pending so a server answering ping in "initializing" state cannot clobber connected=true before the fresh Initialize completes. Includes 30s timeout, structured slog telemetry, concurrent CAS dedup test.

Files: tool_filter.go + tool_filter_test.go (new), session_reset.go + session_reset_test.go (new), manager_connect.go (connectServer/connectViaPool signatures + registerBridgeTools/registerPoolBridgeTools filter logic + reconnPending skip), manager.go (connectAndFilter + connectServer call signature changes), loop_mcp_user.go (filter-at-register block + WithForceReconnect wiring), pool.go (reconnPending skip), bridge_tool.go (session reset detection + WithForceReconnect callback).

* fix(mcp): self-heal grant cache + bypass per-user grant for system/empty userID

Two production fixes for "MCP tool: grant revoked" recurring on song-nhi-v2.

(1) System-user bypass in ListAccessible: Registration uses LoadForAgent(ctx, agentID, "") while execute uses IsAllowed(ctx, agentID, "system", ...). The LEFT JOIN on mcp_user_grants could match a stale disabled row keyed user_id='system' and silently filter the server out only at execute. Skip the join entirely for synthetic owner identities (userID="" or "system") so registration and execute see the same set. Applied to both PostgreSQL (mcp_servers_access.go) and SQLite (mcp_servers_access.go).

(2) Grant-checker no-cache on empty allowByServer: grant_checker.loadEntry now skips the cache write when allowByServer is empty. Without this, a single transient empty result pinned permanent denial until a bus invalidate fired. Re-queries until the empty condition clears, then caches normally. Includes TestStoreGrantChecker_EmptyEntryNotCached.
This commit is contained in:
Duy /zuey/ authored and GitHub committed 2026-05-28 14:04:04 +07:00
1 parent af10f5d708
commit bc3bc25c98
14 files changed
+606 -112

No files matched your search

+25 -1
View File
@@ -188,12 +188,36 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool
// user's BridgeTool (with first user's MCP api_key + pool connection).
// The shared registry holds only shared/non-MCP tools (memory, web,
// exec, …).
//
// Filter tools upfront by the agent's grant (info.ToolAllow / ToolDeny)
// so the LLM never sees tools it cannot call. Without this, every
// per-user MCP server exposes its full tool set and the LLM repeatedly
// triggers the runtime "grant revoked" path (visible in the original
// agent_brain_external screenshot).
hints := mcpbridge.ParseToolHints(srv.Settings)
var filteredOut []string
for _, mcpTool := range entry.MCPTools() {
if !mcpbridge.IsToolAllowed(mcpTool.Name, info.ToolAllow, info.ToolDeny) {
filteredOut = append(filteredOut, mcpTool.Name)
continue
}
bt := mcpbridge.NewBridgeTool(srv.Name, mcpTool, entry.ClientPtr(), srv.ToolPrefix, srv.TimeoutSec, entry.Connected(), srv.ID, l.mcpGrantChecker).
WithHints(hints.Global, hints.HintFor(mcpTool.Name))
WithHints(hints.Global, hints.HintFor(mcpTool.Name)).
WithForceReconnect(entry.RequestForceReconnect())
userTools = append(userTools, bt)
}
if len(filteredOut) > 0 {
slog.Info("mcp.tools.filtered_at_register",
"server", srv.Name,
"server_id", srv.ID,
"user", userID,
"path", "user_cred",
"filtered_count", len(filteredOut),
"filtered_tools", filteredOut,
"allow_size", len(info.ToolAllow),
"deny_size", len(info.ToolDeny),
)
}
}
if len(userTools) > 0 {
+56 -4
View File
@@ -50,6 +50,10 @@ type BridgeTool struct {
timeoutSec int
connected *atomic.Bool
grantChecker GrantChecker // for runtime grant recheck (nil = skip check)
// forceReconnect triggers an out-of-band Initialize when Execute detects
// the server reset its session lifecycle. Optional — nil falls back to
// "connected=false + wait for health loop". Wired via WithForceReconnect.
forceReconnect func(reason string)
}
// NewBridgeTool creates a BridgeTool from an MCP Tool definition.
@@ -120,6 +124,16 @@ func (t *BridgeTool) Description() string {
}
func (t *BridgeTool) Parameters() map[string]any { return t.inputSchema }
// WithForceReconnect attaches a callback used when Execute detects the
// server reset its session lifecycle (see isSessionUninitializedErr).
// The callback should dedupe concurrent invocations internally.
// Optional — leave unset to fall back to flipping connected=false and waiting
// for the health loop's standard reconnect path.
func (t *BridgeTool) WithForceReconnect(fn func(reason string)) *BridgeTool {
t.forceReconnect = fn
return t
}
// WithHints attaches admin-authored description hints to this tool. Hints are
// appended to Description() so the LLM sees server-specific quirks (e.g. "no
// trailing semicolons in code args") without modifying the upstream MCP server.
@@ -175,18 +189,32 @@ func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Re
if t.grantChecker != nil {
agentID := store.AgentIDFromContext(ctx)
userID := store.UserIDFromContext(ctx)
if !t.grantChecker.IsAllowed(ctx, agentID, userID, t.serverID, t.toolName) {
return tools.ErrorResult(fmt.Sprintf("MCP tool %q: grant revoked", t.registeredName))
if allowed, reason := t.grantChecker.IsAllowed(ctx, agentID, userID, t.serverID, t.toolName); !allowed {
// Surface reason + remediation hint so operators can debug from a
// single log line. The tag (server_not_accessible / tool_not_in_allow_list /
// load_failed) maps 1:1 to security.mcp_grant_revoked_at_execute logs.
return tools.ErrorResult(fmt.Sprintf(
"MCP tool %q: grant revoked (reason: %s, server=%s, agent=%s). "+
"Check mcp_agent_grants.enabled / tool_allow for this agent+server, "+
"or re-grant access via dashboard.",
t.registeredName, reason, t.serverName, agentID))
}
}
if !t.connected.Load() {
return tools.ErrorResult(fmt.Sprintf("MCP server %q is disconnected", t.serverName))
return tools.ErrorResult(fmt.Sprintf(
"MCP server %q is disconnected (tool %q). The server may be restarting, "+
"unreachable, or had its credential revoked — check mcp.server.health_failed / "+
"mcp.tool.call.auth_expired logs for the underlying cause.",
t.serverName, t.registeredName))
}
client := t.clientPtr.Load() // atomic load — safe during concurrent reconnect
if client == nil {
return tools.ErrorResult(fmt.Sprintf("MCP server %q has no active client", t.serverName))
return tools.ErrorResult(fmt.Sprintf(
"MCP server %q has no active client (tool %q). Connection was closed "+
"and reconnect has not yet succeeded.",
t.serverName, t.registeredName))
}
callCtx, cancel := context.WithTimeout(ctx, time.Duration(t.timeoutSec)*time.Second)
@@ -234,6 +262,30 @@ func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Re
"latency_ms", latencyMs)
return tools.ErrorResult(fmt.Sprintf("MCP tool %q: credential expired, please retry", t.registeredName))
}
// Detect server-side session reset (FastMCP "invalid during session
// initialization", mcp-go ErrSessionTerminated, etc.). Trigger an
// out-of-band Initialize so the next agent turn finds a healthy
// client; return retry hint to the LLM for the current turn.
if isSessionUninitializedErr(err) {
slog.Warn("mcp.tool.call.session_reset",
"server", t.serverName, "tool", t.registeredName,
"user_id", store.UserIDFromContext(ctx),
"agent_id", store.AgentIDFromContext(ctx),
"latency_ms", latencyMs,
"error", err.Error(),
"action", "force_reconnect_requested")
if t.forceReconnect != nil {
t.forceReconnect("bridge_tool: " + t.registeredName)
} else {
// Best-effort fallback: flip connected so the next pool
// Acquire takes the reconnect branch. Less reliable than
// the explicit force-reconnect path (see WithForceReconnect).
t.connected.Store(false)
}
return tools.ErrorResult(fmt.Sprintf(
"MCP tool %q: server %q reset its session — reconnecting in background, please retry",
t.registeredName, t.serverName))
}
slog.Warn("mcp.tool.call.error",
"server", t.serverName, "tool", t.registeredName,
"user_id", store.UserIDFromContext(ctx),
+32 -11
View File
@@ -14,9 +14,14 @@ import (
// GrantChecker verifies if an agent/user still has access to an MCP server tool.
// Used by BridgeTool.Execute to recheck grants at runtime.
type GrantChecker interface {
// IsAllowed returns true if the agent+user combination has access to the
// specified serverID and toolName. Returns false if grant was revoked.
IsAllowed(ctx context.Context, agentID uuid.UUID, userID string, serverID uuid.UUID, toolName string) bool
// IsAllowed returns (true, "") if the agent+user combination has access
// to the specified serverID and toolName. On denial returns (false, reason)
// where reason is a short machine-friendly tag suitable for logs and for
// surfacing to operators in error messages. Possible reasons:
// - "server_not_accessible" — no agent grant, server disabled, or revoked
// - "tool_not_in_allow_list" — grant exists but tool not in tool_allow
// - "load_failed: <err>" — store query failed (fail-closed)
IsAllowed(ctx context.Context, agentID uuid.UUID, userID string, serverID uuid.UUID, toolName string) (bool, string)
}
// cacheKey identifies a unique agent+user combination for caching.
@@ -63,10 +68,10 @@ func NewStoreGrantChecker(mcpStore store.MCPServerStore, msgBus *bus.MessageBus)
}
// IsAllowed checks if the agent+user has access to the server and tool.
func (gc *storeGrantChecker) IsAllowed(ctx context.Context, agentID uuid.UUID, userID string, serverID uuid.UUID, toolName string) bool {
func (gc *storeGrantChecker) IsAllowed(ctx context.Context, agentID uuid.UUID, userID string, serverID uuid.UUID, toolName string) (bool, string) {
if gc.store == nil {
// No store = config-path mode, skip grant check
return true
return true, ""
}
key := cacheKey{agentID: agentID, userID: userID}
@@ -82,7 +87,22 @@ func (gc *storeGrantChecker) IsAllowed(ctx context.Context, agentID uuid.UUID, u
if err != nil {
slog.Warn("mcp.grant_checker.load_failed", "agent", agentID, "user", userID, "error", err)
// On error, deny access (fail-closed)
return false
return false, "load_failed: " + err.Error()
}
// Self-healing guard: an empty allowByServer means ListAccessible returned
// zero rows. Caching that would pin permanent denial for this (agent,user)
// until a bus invalidate fires — and if the emptiness came from a transient
// condition (race during agent reload, momentary scope misroute, store
// flake that returned nil rows instead of an error), the user stays locked
// out indefinitely. Skipping the cache write here means the next call
// re-queries the store, so the system recovers as soon as the condition
// clears. The cost is one extra ListAccessible per call while truly empty,
// which is acceptable because agents with zero MCP grants don't call MCP
// tools in steady state.
if len(entry.allowByServer) == 0 {
slog.Debug("mcp.grant_checker.empty_no_cache", "agent", agentID, "user", userID)
return gc.checkAccess(entry, serverID, toolName, agentID, userID)
}
gc.cache.Store(key, entry)
@@ -121,7 +141,8 @@ func (gc *storeGrantChecker) loadEntry(ctx context.Context, agentID uuid.UUID, u
}
// checkAccess evaluates if the entry allows access to serverID+toolName.
func (gc *storeGrantChecker) checkAccess(entry *cacheEntry, serverID uuid.UUID, toolName string, agentID uuid.UUID, userID string) bool {
// Returns (allowed, reason). Reason is empty when allowed.
func (gc *storeGrantChecker) checkAccess(entry *cacheEntry, serverID uuid.UUID, toolName string, agentID uuid.UUID, userID string) (bool, string) {
allowSet, hasServer := entry.allowByServer[serverID]
if !hasServer {
slog.Warn("security.mcp_grant_revoked_at_execute",
@@ -131,12 +152,12 @@ func (gc *storeGrantChecker) checkAccess(entry *cacheEntry, serverID uuid.UUID,
"tool", toolName,
"reason", "server_not_accessible",
)
return false
return false, "server_not_accessible"
}
// nil allowSet means "all tools allowed"
if allowSet == nil {
return true
return true, ""
}
// Check if tool is in the allow set
@@ -148,10 +169,10 @@ func (gc *storeGrantChecker) checkAccess(entry *cacheEntry, serverID uuid.UUID,
"tool", toolName,
"reason", "tool_not_in_allow_list",
)
return false
return false, "tool_not_in_allow_list"
}
return true
return true, ""
}
// Invalidate clears the entire cache. Called when grants are modified.
+98 -8
View File
@@ -84,18 +84,18 @@ func TestStoreGrantChecker_CacheHit(t *testing.T) {
ctx := context.Background()
// First call — cache miss, queries store
result1 := gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool")
result1, reason1 := gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool")
if !result1 {
t.Error("expected allowed for accessible server")
t.Errorf("expected allowed for accessible server, reason=%q", reason1)
}
if mockStore.callCount != 1 {
t.Errorf("expected 1 DB call, got %d", mockStore.callCount)
}
// Second call — cache hit, no additional DB call
result2 := gc.IsAllowed(ctx, agentID, userID, serverID, "another_tool")
result2, reason2 := gc.IsAllowed(ctx, agentID, userID, serverID, "another_tool")
if !result2 {
t.Error("expected allowed for accessible server (cached)")
t.Errorf("expected allowed for accessible server (cached), reason=%q", reason2)
}
if mockStore.callCount != 1 {
t.Errorf("expected 1 DB call (cache hit), got %d", mockStore.callCount)
@@ -132,6 +132,96 @@ func TestStoreGrantChecker_InvalidateClearsCache(t *testing.T) {
}
}
func TestStoreGrantChecker_DenyReasonSurfaced(t *testing.T) {
serverID := uuid.New()
agentID := uuid.New()
userID := "test-user"
mockStore := &mockMCPStore{
accessible: []store.MCPAccessInfo{
{
Server: store.MCPServerData{BaseModel: store.BaseModel{ID: serverID}, Name: "test-server"},
ToolAllow: []string{"only_this"},
},
},
}
gc := NewStoreGrantChecker(mockStore, nil)
ctx := context.Background()
// Tool not in allow list — reason must identify why
allowed, reason := gc.IsAllowed(ctx, agentID, userID, serverID, "blocked_tool")
if allowed {
t.Fatal("expected denied for tool not in allow list")
}
if reason != "tool_not_in_allow_list" {
t.Errorf("expected reason=tool_not_in_allow_list, got %q", reason)
}
// Different server entirely — reason: server_not_accessible
otherServer := uuid.New()
allowed, reason = gc.IsAllowed(ctx, agentID, userID, otherServer, "anything")
if allowed {
t.Fatal("expected denied for ungranted server")
}
if reason != "server_not_accessible" {
t.Errorf("expected reason=server_not_accessible, got %q", reason)
}
}
// TestStoreGrantChecker_EmptyEntryNotCached verifies the self-healing guard
// that prevents permanent denial when ListAccessible transiently returns no
// rows. Without this guard, a single empty result pins denial on the (agent,
// user) pair until a bus invalidate event — which can be hours or never.
func TestStoreGrantChecker_EmptyEntryNotCached(t *testing.T) {
serverID := uuid.New()
agentID := uuid.New()
userID := "system"
mockStore := &mockMCPStore{accessible: nil} // simulate transient empty
gc := NewStoreGrantChecker(mockStore, nil)
ctx := context.Background()
// First call hits the store and sees zero rows.
allowed1, reason1 := gc.IsAllowed(ctx, agentID, userID, serverID, "tool")
if allowed1 {
t.Fatalf("expected denied on empty accessible, got allowed (reason=%q)", reason1)
}
if mockStore.callCount != 1 {
t.Fatalf("expected 1 DB call on first miss, got %d", mockStore.callCount)
}
// Second call — if empty had been cached, callCount would stay at 1.
// The guard forces a re-query so the user recovers when the underlying
// transient condition clears.
allowed2, _ := gc.IsAllowed(ctx, agentID, userID, serverID, "tool")
if allowed2 {
t.Fatal("expected denied while accessible is still empty")
}
if mockStore.callCount != 2 {
t.Errorf("expected re-query on empty entry, got callCount=%d", mockStore.callCount)
}
// Flip the store to non-empty (simulates the transient condition clearing).
mockStore.accessible = []store.MCPAccessInfo{
{Server: store.MCPServerData{BaseModel: store.BaseModel{ID: serverID}, Name: "test"}},
}
allowed3, reason3 := gc.IsAllowed(ctx, agentID, userID, serverID, "tool")
if !allowed3 {
t.Errorf("expected allowed after store recovered, got denied (reason=%q)", reason3)
}
if mockStore.callCount != 3 {
t.Errorf("expected fresh query after recovery, got callCount=%d", mockStore.callCount)
}
// After a successful non-empty load, the entry IS cached — subsequent
// calls must NOT re-query (otherwise we lose the cache hit benefit).
gc.IsAllowed(ctx, agentID, userID, serverID, "tool")
if mockStore.callCount != 3 {
t.Errorf("expected cache hit after successful load, got callCount=%d", mockStore.callCount)
}
}
func TestStoreGrantChecker_ToolAllowFilter(t *testing.T) {
serverID := uuid.New()
agentID := uuid.New()
@@ -150,12 +240,12 @@ func TestStoreGrantChecker_ToolAllowFilter(t *testing.T) {
ctx := context.Background()
// Allowed tool
if !gc.IsAllowed(ctx, agentID, userID, serverID, "allowed_tool") {
if allowed, _ := gc.IsAllowed(ctx, agentID, userID, serverID, "allowed_tool"); !allowed {
t.Error("expected true for tool in allow list")
}
// Not allowed tool
if gc.IsAllowed(ctx, agentID, userID, serverID, "blocked_tool") {
if allowed, _ := gc.IsAllowed(ctx, agentID, userID, serverID, "blocked_tool"); allowed {
t.Error("expected false for tool not in allow list")
}
}
@@ -173,7 +263,7 @@ func TestStoreGrantChecker_NoGrant_Denied(t *testing.T) {
ctx := context.Background()
// No grant = denied
if gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool") {
if allowed, _ := gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool"); allowed {
t.Error("expected false when no grant exists")
}
}
@@ -183,7 +273,7 @@ func TestStoreGrantChecker_NilStore_AllowsAll(t *testing.T) {
ctx := context.Background()
// Nil store = config-path mode, skip check
if !gc.IsAllowed(ctx, uuid.New(), "user", uuid.New(), "any_tool") {
if allowed, _ := gc.IsAllowed(ctx, uuid.New(), "user", uuid.New(), "any_tool"); !allowed {
t.Error("expected true when store is nil (config-path mode)")
}
}
+24 -21
View File
@@ -74,6 +74,15 @@ type serverState struct {
reconnAttempts int
healthFailures int // consecutive ping failures (resets on success)
lastErr string
// reconnPending is set when a BridgeTool detects the server reset its
// session lifecycle (FastMCP-style "tools/call invalid during session
// initialization" error) and force-reconnect is in flight. The health
// loop must skip its ping while pending, otherwise a server that still
// answers `ping` in "initializing" state would clobber connected=true
// before the fresh Initialize completes — leaving the pool to keep
// serving the dead session.
reconnPending atomic.Bool
}
// Manager orchestrates MCP server connections and tool registration.
@@ -184,7 +193,9 @@ func (m *Manager) Start(ctx context.Context) error {
continue
}
// Config-path servers have no DB-backed Settings, so no tool hints.
if err := m.connectServer(ctx, name, cfg.Transport, cfg.Command, cfg.Args, cfg.Env, cfg.URL, headers, cfg.ToolPrefix, cfg.TimeoutSec, uuid.Nil, ToolHints{}); err != nil {
// Also no grant-based tool filtering (allow/deny are nil) — config
// trust model is that the operator gates which servers are enabled.
if err := m.connectServer(ctx, name, cfg.Transport, cfg.Command, cfg.Args, cfg.Env, cfg.URL, headers, cfg.ToolPrefix, cfg.TimeoutSec, uuid.Nil, ToolHints{}, nil, nil); err != nil {
slog.Warn("mcp.server.connect_failed", "server", name, "error", err)
errs = append(errs, fmt.Sprintf("%s: %v", name, err))
}
@@ -285,8 +296,10 @@ func (m *Manager) resolveServerCredentials(ctx context.Context, info store.MCPAc
}
}
// connectAndFilter establishes the MCP connection (pool or per-agent mode)
// and applies tool allow/deny filtering from server grants.
// connectAndFilter establishes the MCP connection (pool or per-agent mode).
// Tool allow/deny filtering from server grants is applied upfront inside
// registerBridgeTools / registerPoolBridgeTools so non-allowed tools never
// reach the registry (and thus never reach the LLM).
func (m *Manager) connectAndFilter(ctx context.Context, rs *resolvedServer) error {
srv := rs.info.Server
hints := ParseToolHints(srv.Settings)
@@ -294,25 +307,15 @@ func (m *Manager) connectAndFilter(ctx context.Context, rs *resolvedServer) erro
if m.pool != nil && !rs.hasUserCreds {
// Pool mode: acquire shared connection, create per-agent BridgeTools
tid := store.TenantIDFromContext(ctx)
if err := m.connectViaPool(ctx, tid, srv.Name, srv.Transport, srv.Command,
rs.args, rs.env, srv.URL, rs.headers, srv.ToolPrefix, srv.TimeoutSec, srv.ID, hints); err != nil {
return err
}
} else {
// Per-agent mode: create per-agent connection
if err := m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
rs.args, rs.env, srv.URL, rs.headers,
srv.ToolPrefix, srv.TimeoutSec, srv.ID, hints); err != nil {
return err
}
return m.connectViaPool(ctx, tid, srv.Name, srv.Transport, srv.Command,
rs.args, rs.env, srv.URL, rs.headers, srv.ToolPrefix, srv.TimeoutSec, srv.ID, hints,
rs.info.ToolAllow, rs.info.ToolDeny)
}
// Apply tool filtering from grants
if len(rs.info.ToolAllow) > 0 || len(rs.info.ToolDeny) > 0 {
m.filterTools(srv.Name, rs.info.ToolAllow, rs.info.ToolDeny)
}
return nil
// Per-agent mode: create per-agent connection
return m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
rs.args, rs.env, srv.URL, rs.headers,
srv.ToolPrefix, srv.TimeoutSec, srv.ID, hints,
rs.info.ToolAllow, rs.info.ToolDeny)
}
// LoadForAgent connects MCP servers accessible by a specific agent+user.
+63 -9
View File
@@ -104,14 +104,16 @@ func connectAndDiscover(ctx context.Context, name, transportType, command string
// serverID is the MCP server UUID from DB (uuid.Nil for config-path servers).
// hints carries admin-authored description hints from MCPServerData.Settings.tool_hints;
// pass a zero-value ToolHints{} for config-path servers or when no hints are configured.
func (m *Manager) connectServer(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints) error {
// toolAllow/toolDeny come from the agent's grant — non-allowed tools are filtered out
// upfront so the LLM never sees tools it cannot call. Pass nil for config-path servers.
func (m *Manager) connectServer(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints, toolAllow, toolDeny []string) error {
ss, mcpTools, err := connectAndDiscover(ctx, name, transportType, command, args, env, url, headers, timeoutSec)
if err != nil {
return err
}
// Register tools
registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec, serverID, hints)
// Register tools (filtered by grant's tool_allow/tool_deny upfront)
registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec, serverID, hints, toolAllow, toolDeny)
ss.toolNames = registeredNames
// Create health monitoring context
@@ -143,11 +145,22 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
// registers them in the Manager's registry. Returns registered tool names.
// serverID is the MCP server UUID (uuid.Nil for config-path servers).
// hints.Global applies to all tools; hints.Tools[name] adds a per-tool hint.
func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints) []string {
// toolAllow/toolDeny are evaluated via IsToolAllowed so filtered-out tools
// never get a BridgeTool created — the LLM never sees them, eliminating the
// "registered then runtime-denied" loop that produced repeated grant-revoked
// errors. Pass nil/nil to register every discovered tool.
func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints, toolAllow, toolDeny []string) []string {
var registeredNames []string
var filteredOut []string
for _, mcpTool := range mcpTools {
if !IsToolAllowed(mcpTool.Name, toolAllow, toolDeny) {
filteredOut = append(filteredOut, mcpTool.Name)
continue
}
bt := NewBridgeTool(serverName, mcpTool, &ss.clientPtr, toolPrefix, timeoutSec, &ss.connected, serverID, m.grantChecker).
WithHints(hints.Global, hints.HintFor(mcpTool.Name))
WithHints(hints.Global, hints.HintFor(mcpTool.Name)).
WithForceReconnect(func(reason string) { ss.requestForceReconnect(reason) })
if _, exists := m.registry.Get(bt.Name()); exists {
slog.Warn("mcp.tool.name_collision",
@@ -161,6 +174,17 @@ func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, se
m.registry.Register(bt)
registeredNames = append(registeredNames, bt.Name())
}
if len(filteredOut) > 0 {
slog.Info("mcp.tools.filtered_at_register",
"server", serverName,
"server_id", serverID,
"path", "manager",
"filtered_count", len(filteredOut),
"filtered_tools", filteredOut,
"allow_size", len(toolAllow),
"deny_size", len(toolDeny),
)
}
return registeredNames
}
@@ -168,14 +192,16 @@ func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, se
// per-agent BridgeTools pointing to the shared client/connected pointers.
// serverID is the MCP server UUID from DB. hints carries admin-authored
// description hints from MCPServerData.Settings.tool_hints.
func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints) error {
// toolAllow/toolDeny come from the agent's grant — non-allowed tools are filtered
// out upfront so the LLM never sees tools it cannot call.
func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints, toolAllow, toolDeny []string) error {
entry, err := m.pool.Acquire(ctx, tenantID, name, transportType, command, args, env, url, headers, timeoutSec)
if err != nil {
return err
}
// Create per-agent BridgeTools from the pool's shared connection
registeredNames := m.registerPoolBridgeTools(entry, name, toolPrefix, timeoutSec, serverID, hints)
registeredNames := m.registerPoolBridgeTools(entry, name, toolPrefix, timeoutSec, serverID, hints, toolAllow, toolDeny)
// Track server state and per-agent tool names.
// poolServers/poolToolNames keyed by plain name for Close() iteration.
@@ -214,11 +240,20 @@ func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name,
// pointing to the shared client/connected pointers. Returns registered tool names.
// serverID is the MCP server UUID from DB.
// hints.Global applies to all tools; hints.Tools[name] adds a per-tool hint.
func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints) []string {
// toolAllow/toolDeny are evaluated via IsToolAllowed so filtered-out tools never
// get a BridgeTool created (LLM never sees them).
func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID, hints ToolHints, toolAllow, toolDeny []string) []string {
var registeredNames []string
var filteredOut []string
for _, mcpTool := range entry.tools {
if !IsToolAllowed(mcpTool.Name, toolAllow, toolDeny) {
filteredOut = append(filteredOut, mcpTool.Name)
continue
}
bt := NewBridgeTool(serverName, mcpTool, &entry.state.clientPtr, toolPrefix, timeoutSec, &entry.state.connected, serverID, m.grantChecker).
WithHints(hints.Global, hints.HintFor(mcpTool.Name))
WithHints(hints.Global, hints.HintFor(mcpTool.Name)).
WithForceReconnect(entry.RequestForceReconnect())
if _, exists := m.registry.Get(bt.Name()); exists {
slog.Warn("mcp.tool.name_collision",
@@ -232,6 +267,17 @@ func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPref
m.registry.Register(bt)
registeredNames = append(registeredNames, bt.Name())
}
if len(filteredOut) > 0 {
slog.Info("mcp.tools.filtered_at_register",
"server", serverName,
"server_id", serverID,
"path", "pool",
"filtered_count", len(filteredOut),
"filtered_tools", filteredOut,
"allow_size", len(toolAllow),
"deny_size", len(toolDeny),
)
}
return registeredNames
}
@@ -283,6 +329,14 @@ func (m *Manager) healthLoop(ctx context.Context, ss *serverState) {
case <-ctx.Done():
return
case <-ticker.C:
// Skip ping while a force-reconnect is in flight. Some servers
// answer `ping` in the post-reset "initializing" state; pinging
// here would race the recovery goroutine and clobber
// connected=true before the fresh Initialize completes.
if ss.reconnPending.Load() {
slog.Debug("mcp.server.health_skip", "server", ss.name, "reason", "reconnect_pending")
continue
}
if err := ss.client.Ping(ctx); err != nil {
if isMethodNotFound(err) {
ss.connected.Store(true)
-57
View File
@@ -149,60 +149,3 @@ func DiscoverTools(ctx context.Context, transportType, command string, args []st
}
return result, nil
}
// filterTools removes tools from the registry that don't match the allow/deny lists.
func (m *Manager) filterTools(serverName string, allow, deny []string) {
m.mu.Lock()
defer m.mu.Unlock()
// Get the tool names list (pool-backed or standalone)
var toolNames []string
_, isPool := m.poolServers[serverName]
if isPool {
toolNames = m.poolToolNames[serverName]
} else if ss, ok := m.servers[serverName]; ok {
toolNames = ss.toolNames
} else {
return
}
allowSet := toSet(allow)
denySet := toSet(deny)
var kept []string
for _, toolName := range toolNames {
bt, ok := m.registry.Get(toolName)
if !ok {
continue
}
bridge, ok := bt.(*BridgeTool)
if !ok {
kept = append(kept, toolName)
continue
}
origName := bridge.OriginalName()
// Deny takes priority
if _, denied := denySet[origName]; denied {
m.registry.Unregister(toolName)
continue
}
// If allow list is set, only keep tools in the allow list
if len(allowSet) > 0 {
if _, allowed := allowSet[origName]; !allowed {
m.registry.Unregister(toolName)
continue
}
}
kept = append(kept, toolName)
}
// Update the correct tool names list
if isPool {
m.poolToolNames[serverName] = kept
} else {
m.servers[serverName].toolNames = kept
}
}
+17
View File
@@ -598,6 +598,15 @@ func (e *poolEntry) Connected() *atomic.Bool { return &e.state.connected }
// MCPTools returns the discovered MCP tool definitions for this pool entry.
func (e *poolEntry) MCPTools() []mcpgo.Tool { return e.tools }
// RequestForceReconnect triggers an out-of-band Initialize when a BridgeTool
// detects the server-side session was reset (see isSessionUninitializedErr).
// Returns a closure to keep BridgeTool decoupled from *serverState.
// Concurrent invocations dedupe via the underlying CAS guard.
func (e *poolEntry) RequestForceReconnect() func(reason string) {
ss := e.state
return func(reason string) { ss.requestForceReconnect(reason) }
}
// poolHealthLoop is a standalone health loop for pool-managed connections.
// After consecutive ping failures, it attempts a full reconnect by creating
// a fresh client, mirroring the Manager.tryReconnect slow path.
@@ -610,6 +619,14 @@ func poolHealthLoop(ctx context.Context, ss *serverState) {
case <-ctx.Done():
return
case <-ticker.C:
// Skip ping while a force-reconnect is in flight (see
// requestForceReconnect). Pinging here races the recovery
// goroutine and would clobber connected=true on servers that
// answer `ping` even in the post-reset "initializing" state.
if ss.reconnPending.Load() {
slog.Debug("mcp.pool.health_skip", "server", ss.name, "reason", "reconnect_pending")
continue
}
if err := ss.client.Ping(ctx); err != nil {
if isMethodNotFound(err) {
ss.connected.Store(true)
+84
View File
@@ -0,0 +1,84 @@
package mcp
import (
"context"
"log/slog"
"strings"
"time"
)
// forceReconnectTimeout bounds the synchronous Initialize during a
// session-reset recovery so a slow server cannot wedge the goroutine.
const forceReconnectTimeout = 30 * time.Second
// isSessionUninitializedErr detects server-side responses that indicate the
// MCP session lifecycle was reset and the cached client must re-Initialize
// before further tool calls succeed.
//
// Background: FastMCP / mcp-spec stateful HTTP servers track an "initializing
// → initialized" state per Mcp-Session-Id. When the server restarts, GCs idle
// sessions, or scales down, the in-memory state vanishes; the client keeps
// reusing the same SID and the server treats it as a fresh "initializing"
// session that rejects `tools/call` until `notifications/initialized` arrives.
//
// The mcp-go transport only maps HTTP 404 → ErrSessionTerminated. FastMCP-
// style servers return HTTP 200 with a JSON-RPC error body instead, which
// surfaces here as a plain Go error. We string-match because there is no
// dedicated error code in the JSON-RPC spec for this lifecycle violation.
//
// Known phrasings (extend as new servers surface):
// - "method <X> is invalid during session initialization" (FastMCP / Python mcp)
// - "session not initialized" (mcp-go server, some Node implementations)
// - "session terminated" (catch-all for mcp-go ErrSessionTerminated text)
func isSessionUninitializedErr(err error) bool {
if err == nil {
return false
}
msg := strings.ToLower(err.Error())
return strings.Contains(msg, "invalid during session initialization") ||
strings.Contains(msg, "session not initialized") ||
strings.Contains(msg, "session terminated")
}
// requestForceReconnect kicks off a fresh Initialize handshake out-of-band
// because a BridgeTool detected the server lost session state. Concurrent
// calls are deduped via reconnPending CAS so N failing tool calls in flight
// trigger exactly one reconnect.
//
// Runs asynchronously: the caller (BridgeTool.Execute) returns the original
// error to the agent loop without waiting. By the next tool call attempt the
// fresh client will be in place via the atomic clientPtr swap inside
// fullReconnect.
func (ss *serverState) requestForceReconnect(reason string) {
if !ss.reconnPending.CompareAndSwap(false, true) {
slog.Debug("mcp.session_reset.dedup",
"server", ss.name, "reason", reason)
return
}
slog.Warn("mcp.session_reset.detected",
"server", ss.name,
"transport", ss.transport,
"reason", reason,
"action", "force_reconnect")
go func() {
defer ss.reconnPending.Store(false)
ctx, cancel := context.WithTimeout(context.Background(), forceReconnectTimeout)
defer cancel()
ss.connected.Store(false)
start := time.Now()
if fullReconnect(ctx, ss) {
slog.Info("mcp.session_reset.recovered",
"server", ss.name,
"latency_ms", time.Since(start).Milliseconds())
return
}
slog.Warn("mcp.session_reset.recovery_failed",
"server", ss.name,
"latency_ms", time.Since(start).Milliseconds(),
"hint", "next health-loop tick will retry via standard reconnectWithBackoff")
}()
}
+88
View File
@@ -0,0 +1,88 @@
package mcp
import (
"errors"
"sync"
"sync/atomic"
"testing"
)
func TestIsSessionUninitializedErr(t *testing.T) {
cases := []struct {
name string
err error
want bool
}{
{"nil", nil, false},
{"unrelated", errors.New("connection refused"), false},
// FastMCP / Python mcp server phrasing — the exact text observed in
// the production trace that motivated this detector.
{"fastmcp tools/call", errors.New(`method "tools/call" is invalid during session initialization`), true},
{"fastmcp other method", errors.New(`method "resources/list" is invalid during session initialization`), true},
// mcp-go server / Node implementations.
{"session not initialized", errors.New("session not initialized"), true},
// mcp-go transport ErrSessionTerminated text (HTTP 404 path).
{"session terminated", errors.New("session terminated (404). need to re-initialize"), true},
// Case-insensitive matching keeps detection robust to upstream
// rewording without dragging the maintainer into a string-match
// audit every time a server logs slightly differently.
{"mixed case", errors.New("Session Not Initialized"), true},
// 401 must NOT match — handled by isUnauthorizedErr to drive a
// different recovery path (credential purge).
{"unauthorized 401", errors.New("unauthorized (401)"), false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := isSessionUninitializedErr(tc.err); got != tc.want {
t.Errorf("isSessionUninitializedErr(%v) = %v, want %v", tc.err, got, tc.want)
}
})
}
}
// TestRequestForceReconnect_Dedup verifies the CAS guard collapses N
// concurrent failing tool calls into a single reconnect in flight. Without
// this, a high-QPS agent burst would launch one fullReconnect per failed
// call — defeating the purpose of pooling and hammering the recovering
// server during a restart window.
func TestRequestForceReconnect_Dedup(t *testing.T) {
// Stand up a serverState whose Pending flag starts cleared. We don't
// run the goroutine body to completion (that requires a real client);
// the CAS check itself is what we're verifying. The goroutine will
// fail fast on the nil client and clear the flag — but only after
// the test observation point.
ss := &serverState{name: "test", transport: "streamable-http"}
var attempted atomic.Int32
var wg sync.WaitGroup
const concurrent = 20
for range concurrent {
wg.Add(1)
go func() {
defer wg.Done()
// Inline the CAS half of requestForceReconnect — testing the
// real method would race the cleanup goroutine. The dedup
// guarantee lives entirely in the CAS, so this is faithful.
if ss.reconnPending.CompareAndSwap(false, true) {
attempted.Add(1)
}
}()
}
wg.Wait()
if got := attempted.Load(); got != 1 {
t.Errorf("expected exactly 1 reconnect attempt after %d concurrent requests, got %d", concurrent, got)
}
if !ss.reconnPending.Load() {
t.Error("expected reconnPending to be set after first CAS")
}
// Once cleared (simulating reconnect completion), the next request must
// be allowed through — otherwise a single transient reset would freeze
// recovery forever.
ss.reconnPending.Store(false)
if !ss.reconnPending.CompareAndSwap(false, true) {
t.Error("expected CAS to succeed after pending flag cleared")
}
}
+27
View File
@@ -0,0 +1,27 @@
package mcp
// IsToolAllowed evaluates whether a tool (by original MCP name) passes the
// allow/deny filter from a grant. Used at registration time so non-allowed
// tools never reach the LLM, and at runtime by the grant checker.
//
// Semantics (matching ListAccessible / GrantChecker.checkAccess):
// - Empty allow + empty deny → allowed (no filter configured)
// - Tool in deny → denied (deny takes priority)
// - Non-empty allow + tool NOT in allow → denied
// - Otherwise → allowed
func IsToolAllowed(toolName string, allow, deny []string) bool {
for _, t := range deny {
if t == toolName {
return false
}
}
if len(allow) == 0 {
return true
}
for _, t := range allow {
if t == toolName {
return true
}
}
return false
}
+31
View File
@@ -0,0 +1,31 @@
package mcp
import "testing"
func TestIsToolAllowed(t *testing.T) {
cases := []struct {
name string
tool string
allow []string
deny []string
expected bool
}{
{"no filter allows everything", "search", nil, nil, true},
{"empty allow allows everything", "search", []string{}, []string{}, true},
{"tool in allow", "search", []string{"search", "list"}, nil, true},
{"tool not in allow", "delete", []string{"search", "list"}, nil, false},
{"tool in deny", "delete", nil, []string{"delete"}, false},
{"deny wins over allow", "delete", []string{"delete"}, []string{"delete"}, false},
{"allow + deny coexist (allowed)", "search", []string{"search"}, []string{"delete"}, true},
{"allow + deny coexist (denied by allow)", "list", []string{"search"}, []string{"delete"}, false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := IsToolAllowed(tc.tool, tc.allow, tc.deny)
if got != tc.expected {
t.Errorf("IsToolAllowed(%q, allow=%v, deny=%v) = %v, want %v",
tc.tool, tc.allow, tc.deny, got, tc.expected)
}
})
}
}
+35 -1
View File
@@ -2,6 +2,7 @@ package pg
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"time"
@@ -148,7 +149,33 @@ func (s *PGMCPServerStore) RevokeFromUser(ctx context.Context, serverID uuid.UUI
// --- Resolution ---
func (s *PGMCPServerStore) ListAccessible(ctx context.Context, agentID uuid.UUID, userID string) ([]store.MCPAccessInfo, error) {
tClause, tArgs, _, err := scopeClauseAlias(ctx, 3, "ms")
tClause, tArgs, _, err := scopeClauseAlias(ctx, 2, "ms")
if err != nil {
return nil, err
}
// Synthetic owner identities ("" at startup registration, "system" for
// WS direct/owner chat) are not real external actors — they cannot have
// a per-user grant that anyone deliberately provisioned. Skipping the
// mcp_user_grants join for those keeps registration (userID="") and
// execute (userID="system") symmetric. Without this, a stale row in
// mcp_user_grants with user_id='system' AND enabled=false would silently
// hide tools at execute time that registration saw fine, producing
// "grant revoked (reason: server_not_accessible)" with no DB change.
if userID == "" || userID == "system" {
rows, err := s.db.QueryContext(ctx,
`SELECT ms.id, ms.name, ms.display_name, ms.transport, ms.command, ms.args, ms.url, ms.headers, ms.env,
ms.api_key, ms.tool_prefix, ms.timeout_sec, ms.settings, ms.enabled, ms.created_by, ms.created_at, ms.updated_at,
mag.tool_allow, mag.tool_deny
FROM mcp_servers ms
INNER JOIN mcp_agent_grants mag ON ms.id = mag.server_id AND mag.agent_id = $1 AND mag.enabled = true
WHERE ms.enabled = true`+tClause,
append([]any{agentID}, tArgs...)...)
if err != nil {
return nil, err
}
return s.scanAccessibleRows(rows)
}
tClause, tArgs, _, err = scopeClauseAlias(ctx, 3, "ms")
if err != nil {
return nil, err
}
@@ -165,6 +192,13 @@ func (s *PGMCPServerStore) ListAccessible(ctx context.Context, agentID uuid.UUID
if err != nil {
return nil, err
}
return s.scanAccessibleRows(rows)
}
// scanAccessibleRows decodes the shared SELECT projection used by both the
// system-user (no per-user-grant filter) and external-user paths of
// ListAccessible. Kept private so callers go through ListAccessible.
func (s *PGMCPServerStore) scanAccessibleRows(rows *sql.Rows) ([]store.MCPAccessInfo, error) {
defer rows.Close()
result := make([]store.MCPAccessInfo, 0)
@@ -4,6 +4,7 @@ package sqlitestore
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"log/slog"
@@ -194,6 +195,24 @@ func (s *SQLiteMCPServerStore) ListAccessible(ctx context.Context, agentID uuid.
if err != nil {
return nil, err
}
// Symmetric with PG ListAccessible: synthetic owner identities ("" at
// registration, "system" for WS direct chat) skip the per-user-grant join
// so a stale disabled user_grants row never silently hides a server that
// agent_grants enables. See internal/store/pg/mcp_servers_access.go.
if userID == "" || userID == "system" {
rows, err := s.db.QueryContext(ctx,
`SELECT ms.id, ms.name, ms.display_name, ms.transport, ms.command, ms.args, ms.url, ms.headers, ms.env,
ms.api_key, ms.tool_prefix, ms.timeout_sec, ms.settings, ms.enabled, ms.created_by, ms.created_at, ms.updated_at,
mag.tool_allow, mag.tool_deny
FROM mcp_servers ms
INNER JOIN mcp_agent_grants mag ON ms.id = mag.server_id AND mag.agent_id = ? AND mag.enabled = 1
WHERE ms.enabled = 1`+tClause,
append([]any{agentID}, tArgs...)...)
if err != nil {
return nil, err
}
return s.scanAccessibleRows(rows)
}
rows, err := s.db.QueryContext(ctx,
`SELECT ms.id, ms.name, ms.display_name, ms.transport, ms.command, ms.args, ms.url, ms.headers, ms.env,
ms.api_key, ms.tool_prefix, ms.timeout_sec, ms.settings, ms.enabled, ms.created_by, ms.created_at, ms.updated_at,
@@ -207,6 +226,13 @@ func (s *SQLiteMCPServerStore) ListAccessible(ctx context.Context, agentID uuid.
if err != nil {
return nil, err
}
return s.scanAccessibleRows(rows)
}
// scanAccessibleRows decodes the shared SELECT projection used by both the
// system-user (no per-user-grant filter) and external-user paths of
// ListAccessible.
func (s *SQLiteMCPServerStore) scanAccessibleRows(rows *sql.Rows) ([]store.MCPAccessInfo, error) {
defer rows.Close()
result := make([]store.MCPAccessInfo, 0)