mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 06:17:41 +00:00
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:
1 parent
af10f5d708
commit
bc3bc25c98
14 files changed
+606
-112
No files matched your search
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}()
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user