From 8b8da3a3dd2f56e62859f8a3063ddaa57a0e9d2b Mon Sep 17 00:00:00 2001 From: viettranx Date: Thu, 16 Apr 2026 17:07:56 +0700 Subject: [PATCH] fix(mcp): unified grant revocation with per-agent tool group isolation - Add GrantChecker interface with cache + event-bus invalidation for revalidating MCP tool grants at execution time (not just dispatch) - Move toolGroups from package-level global to per-Registry field; Clone() gives each agent Loop isolated tool groups, eliminating cross-agent race condition - BridgeTool.Execute rechecks grant validity before executing MCP methods; uses serverID field for cache key - Wire MCPGrantChecker through ResolverDeps to MCP Manager and LoopConfig - Add integration test for grant revocation scenarios - Update policy functions to accept *Registry parameter --- cmd/gateway_managed.go | 3 + internal/agent/loop_lazy_mcp_test.go | 8 +- internal/agent/loop_mcp_user.go | 4 +- internal/agent/loop_types.go | 19 +- internal/agent/resolver.go | 7 + internal/http/skills_import.go | 14 ++ internal/http/skills_upload.go | 14 ++ internal/mcp/bridge_tool.go | 22 +- internal/mcp/bridge_tool_test.go | 9 +- internal/mcp/grant_checker.go | 160 +++++++++++++ internal/mcp/grant_checker_test.go | 189 +++++++++++++++ internal/mcp/manager.go | 24 +- internal/mcp/manager_connect.go | 25 +- internal/mcp/manager_lazy_test.go | 3 +- internal/mcp/manager_tools.go | 9 +- internal/tools/policy.go | 187 ++++++++------- internal/tools/policy_mcp_test.go | 25 +- internal/tools/policy_race_test.go | 171 ++++++++++++++ internal/tools/registry.go | 124 +++++++++- tests/integration/mcp_grant_revoke_test.go | 261 +++++++++++++++++++++ 20 files changed, 1128 insertions(+), 150 deletions(-) create mode 100644 internal/mcp/grant_checker.go create mode 100644 internal/mcp/grant_checker_test.go create mode 100644 internal/tools/policy_race_test.go create mode 100644 tests/integration/mcp_grant_revoke_test.go diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index 3d1e3edf..af837ef5 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -127,8 +127,10 @@ func wireExtras( // 5. Shared MCP connection pool (eliminates duplicate connections across agents) var mcpPool *mcpbridge.Pool + var mcpGrantChecker mcpbridge.GrantChecker if stores.MCP != nil { mcpPool = mcpbridge.NewPool(mcpbridge.DefaultPoolConfig()) + mcpGrantChecker = mcpbridge.NewStoreGrantChecker(stores.MCP, msgBus) } // 6. Set up agent resolver: lazy-creates Loops from DB @@ -221,6 +223,7 @@ func wireExtras( BuiltinToolStore: stores.BuiltinTools, MCPStore: stores.MCP, MCPPool: mcpPool, + MCPGrantChecker: mcpGrantChecker, ConfigPermStore: stores.ConfigPermissions, MediaStore: mediaStore, ModelPricing: appCfg.Telemetry.ModelPricing, diff --git a/internal/agent/loop_lazy_mcp_test.go b/internal/agent/loop_lazy_mcp_test.go index 1b3d9da5..c5deea7c 100644 --- a/internal/agent/loop_lazy_mcp_test.go +++ b/internal/agent/loop_lazy_mcp_test.go @@ -337,13 +337,15 @@ func TestLoop_LazyMCP_DenyList_GroupDeny(t *testing.T) { return true }) - // Register a custom group containing the MCP tool. - tools.RegisterToolGroup("mcp_test", []string{"mcp_svc__get_data"}) - defer tools.UnregisterToolGroup("mcp_test") + // Register a custom group containing the MCP tool (using per-Registry method). + reg.RegisterToolGroup("mcp_test", []string{"mcp_svc__get_data"}) + defer reg.UnregisterToolGroup("mcp_test") pe := tools.NewPolicyEngine(&config.ToolsConfig{ Deny: []string{"group:mcp_test"}, }) + // PolicyEngine needs registry to expand group:mcp_test + pe.SetRegistry(reg) if !pe.IsDenied("mcp_svc__get_data", nil) { t.Error("tool should be denied via group: pattern") diff --git a/internal/agent/loop_mcp_user.go b/internal/agent/loop_mcp_user.go index dd842c9d..83030785 100644 --- a/internal/agent/loop_mcp_user.go +++ b/internal/agent/loop_mcp_user.go @@ -89,7 +89,7 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool // shared tool registry so ExecuteWithContext can resolve them by name. reg, _ := l.tools.(*tools.Registry) for _, mcpTool := range entry.MCPTools() { - bt := mcpbridge.NewBridgeTool(srv.Name, mcpTool, entry.ClientPtr(), srv.ToolPrefix, srv.TimeoutSec, entry.Connected()) + bt := mcpbridge.NewBridgeTool(srv.Name, mcpTool, entry.ClientPtr(), srv.ToolPrefix, srv.TimeoutSec, entry.Connected(), srv.ID, l.mcpGrantChecker) // Register in registry so ExecuteWithContext can find them. // Skip if already registered (another user loaded this server with same tool names). if reg != nil { @@ -109,7 +109,7 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool for _, t := range userTools { names = append(names, t.Name()) } - tools.MergeToolGroup("mcp", names) + l.registry.MergeToolGroup("mcp", names) slog.Info("mcp.user_tools_loaded", "user", userID, "tools", len(userTools)) } return userTools diff --git a/internal/agent/loop_types.go b/internal/agent/loop_types.go index 8512c882..7a2d5743 100644 --- a/internal/agent/loop_types.go +++ b/internal/agent/loop_types.go @@ -112,6 +112,7 @@ type Loop struct { domainBus eventbus.DomainEventBus // V3 domain event bus for consolidation pipeline sessions store.SessionStore tools tools.ToolExecutor + registry *tools.Registry // direct registry access for MergeToolGroup (per-Registry tool groups) toolPolicy *tools.PolicyEngine // optional: filters tools sent to LLM agentToolPolicy *config.ToolPolicySpec // per-agent tool policy from DB (nil = no restrictions) activeRuns atomic.Int32 // number of currently executing runs @@ -136,10 +137,11 @@ type Loop struct { userSetups sync.Map // userID → *userSetup (workspace + seeding state, per Loop instance) // Per-user MCP tools: servers requiring user credentials get connected per-request. - mcpStore store.MCPServerStore // for credential lookup - mcpPool *mcpbridge.Pool // user-keyed connection pool - mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds - mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools) + mcpStore store.MCPServerStore // for credential lookup + mcpPool *mcpbridge.Pool // user-keyed connection pool + mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds + mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools) + mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) // Compaction config (memory flush settings) compactionCfg *config.CompactionConfig @@ -416,9 +418,10 @@ type LoopConfig struct { MemoryStore store.MemoryStore // Per-user MCP tools (servers requiring per-user credentials) - MCPStore store.MCPServerStore // for credential lookup - MCPPool *mcpbridge.Pool // user-keyed connection pool - MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds + MCPStore store.MCPServerStore // for credential lookup + MCPPool *mcpbridge.Pool // user-keyed connection pool + MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds + MCPGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip) // V3 orchestration mode (resolved by resolver, controls tool visibility) OrchMode OrchestrationMode @@ -491,6 +494,7 @@ func NewLoop(cfg LoopConfig) *Loop { hookDispatcher: cfg.HookDispatcher, sessions: cfg.Sessions, tools: cfg.Tools, + registry: cfg.Tools, toolPolicy: cfg.ToolPolicy, agentToolPolicy: cfg.AgentToolPolicy, onEvent: cfg.OnEvent, @@ -540,6 +544,7 @@ func NewLoop(cfg LoopConfig) *Loop { mcpStore: cfg.MCPStore, mcpPool: cfg.MCPPool, mcpUserCredSrvs: cfg.MCPUserCredSrvs, + mcpGrantChecker: cfg.MCPGrantChecker, orchMode: cfg.OrchMode, delegateTargets: cfg.DelegateTargets, evolutionMetricsStore: cfg.EvolutionMetricsStore, diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index 66dd3b92..eec9dfd9 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -80,6 +80,9 @@ type ResolverDeps struct { // Shared MCP connection pool — eliminates duplicate connections across agents MCPPool *mcpbridge.Pool + // MCP grant checker — for runtime grant verification at BridgeTool.Execute + MCPGrantChecker mcpbridge.GrantChecker + // Skill access store — for per-agent skill visibility filtering SkillAccessStore store.SkillAccessStore @@ -304,6 +307,9 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { if deps.MCPPool != nil { mcpOpts = append(mcpOpts, mcpbridge.WithPool(deps.MCPPool)) } + if deps.MCPGrantChecker != nil { + mcpOpts = append(mcpOpts, mcpbridge.WithGrantChecker(deps.MCPGrantChecker)) + } mcpMgr := mcpbridge.NewManager(toolsReg, mcpOpts...) if err := mcpMgr.LoadForAgent(ctx, ag.ID, ""); err != nil { slog.Warn("failed to load MCP servers for agent", "agent", agentKey, "error", err) @@ -521,6 +527,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc { MCPStore: deps.MCPStore, MCPPool: deps.MCPPool, MCPUserCredSrvs: mcpUserCredSrvs, + MCPGrantChecker: deps.MCPGrantChecker, OrchMode: orchMode, DelegateTargets: delegateTargets, EvolutionMetricsStore: evoMetricsStore, diff --git a/internal/http/skills_import.go b/internal/http/skills_import.go index f7022469..57744701 100644 --- a/internal/http/skills_import.go +++ b/internal/http/skills_import.go @@ -15,6 +15,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/config" "github.com/nextlevelbuilder/goclaw/internal/i18n" + "github.com/nextlevelbuilder/goclaw/internal/skills" "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/store/pg" ) @@ -132,6 +133,19 @@ func (h *SkillsHandler) doSkillsImport(ctx context.Context, r io.Reader, userID continue } + // Security guard: scan SKILL.md for malicious content BEFORE any disk/DB write + if entry.skillMD != nil { + violations, safe := skills.GuardSkillContent(string(entry.skillMD)) + if !safe { + slog.Warn("security.skills.import_rejected", + "slug", slug, + "violations", len(violations), + "first_rule", violations[0].Reason) + summary.SkillsSkipped++ + continue + } + } + var meta struct { Name string `json:"name"` Slug string `json:"slug"` diff --git a/internal/http/skills_upload.go b/internal/http/skills_upload.go index b7af06bb..fbbbb6de 100644 --- a/internal/http/skills_upload.go +++ b/internal/http/skills_upload.go @@ -106,6 +106,20 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) { return } + // Security guard: scan for malicious content BEFORE any disk/DB write + violations, safe := skills.GuardSkillContent(skillContent) + if !safe { + slog.Warn("security.skills.upload_rejected", + "user_id", userID, + "violations", len(violations), + "first_rule", violations[0].Reason) + writeJSON(w, http.StatusBadRequest, map[string]any{ + "error": i18n.T(locale, i18n.MsgInvalidRequest, "skill content failed security scan"), + "violations": skills.FormatGuardViolations(violations), + }) + return + } + name, description, slug, frontmatter := skills.ParseSkillFrontmatter(skillContent) if name == "" { writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgRequired, "name in SKILL.md frontmatter")}) diff --git a/internal/mcp/bridge_tool.go b/internal/mcp/bridge_tool.go index 69026a86..9bcaf338 100644 --- a/internal/mcp/bridge_tool.go +++ b/internal/mcp/bridge_tool.go @@ -8,8 +8,10 @@ import ( "sync/atomic" "time" + "github.com/google/uuid" mcpclient "github.com/mark3labs/mcp-go/client" mcpgo "github.com/mark3labs/mcp-go/mcp" + "github.com/nextlevelbuilder/goclaw/internal/store" "github.com/nextlevelbuilder/goclaw/internal/tools" ) @@ -19,14 +21,16 @@ import ( // safe reconnection without data races. type BridgeTool struct { serverName string - toolName string // original MCP tool name - registeredName string // may include prefix: "{prefix}__{toolName}" + serverID uuid.UUID // MCP server ID (for grant recheck) + toolName string // original MCP tool name + registeredName string // may include prefix: "{prefix}__{toolName}" description string inputSchema map[string]any // JSON Schema for parameters requiredSet map[string]bool clientPtr *atomic.Pointer[mcpclient.Client] // shared with serverState for atomic swap on reconnect timeoutSec int connected *atomic.Bool + grantChecker GrantChecker // for runtime grant recheck (nil = skip check) } // NewBridgeTool creates a BridgeTool from an MCP Tool definition. @@ -34,7 +38,8 @@ type BridgeTool struct { // If prefix is empty, it is auto-derived from the server name. // clientPtr is a shared atomic pointer from serverState — reconnection swaps it // atomically, and all BridgeTools see the new client without explicit notification. -func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Pointer[mcpclient.Client], prefix string, timeoutSec int, connected *atomic.Bool) *BridgeTool { +// serverID and grantChecker are optional — pass uuid.Nil and nil for config-path mode. +func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Pointer[mcpclient.Client], prefix string, timeoutSec int, connected *atomic.Bool, serverID uuid.UUID, grantChecker GrantChecker) *BridgeTool { name := mcpTool.Name effectivePrefix := ensureMCPPrefix(prefix, serverName) registered := effectivePrefix + "__" + name @@ -52,6 +57,7 @@ func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Poin return &BridgeTool{ serverName: serverName, + serverID: serverID, toolName: name, registeredName: registered, description: mcpTool.Description, @@ -60,6 +66,7 @@ func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Poin clientPtr: clientPtr, timeoutSec: timeoutSec, connected: connected, + grantChecker: grantChecker, } } @@ -99,6 +106,15 @@ func (t *BridgeTool) OriginalName() string { return t.toolName } func (t *BridgeTool) IsConnected() bool { return t.connected.Load() } func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Result { + // Recheck grant before execution — defense against revoked grants + 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 !t.connected.Load() { return tools.ErrorResult(fmt.Sprintf("MCP server %q is disconnected", t.serverName)) } diff --git a/internal/mcp/bridge_tool_test.go b/internal/mcp/bridge_tool_test.go index cd7d751e..7846fd80 100644 --- a/internal/mcp/bridge_tool_test.go +++ b/internal/mcp/bridge_tool_test.go @@ -3,6 +3,7 @@ package mcp import ( "testing" + "github.com/google/uuid" mcpgo "github.com/mark3labs/mcp-go/mcp" ) @@ -93,7 +94,7 @@ func TestBridgeToolNaming(t *testing.T) { } // Without prefix → auto-derived from server name - bt := NewBridgeTool("myserver", mcpTool, nil, "", 30, nil) + bt := NewBridgeTool("myserver", mcpTool, nil, "", 30, nil, uuid.Nil, nil) if bt.Name() != "mcp_myserver__query" { t.Errorf("expected name=mcp_myserver__query, got %s", bt.Name()) } @@ -105,7 +106,7 @@ func TestBridgeToolNaming(t *testing.T) { } // With non-mcp_ prefix → gets mcp_ prepended - bt2 := NewBridgeTool("myserver", mcpTool, nil, "pg", 0, nil) + bt2 := NewBridgeTool("myserver", mcpTool, nil, "pg", 0, nil, uuid.Nil, nil) if bt2.Name() != "mcp_pg__query" { t.Errorf("expected name=mcp_pg__query, got %s", bt2.Name()) } @@ -114,13 +115,13 @@ func TestBridgeToolNaming(t *testing.T) { } // With mcp_ prefix → unchanged - bt3 := NewBridgeTool("myserver", mcpTool, nil, "mcp_pg", 0, nil) + bt3 := NewBridgeTool("myserver", mcpTool, nil, "mcp_pg", 0, nil, uuid.Nil, nil) if bt3.Name() != "mcp_pg__query" { t.Errorf("expected name=mcp_pg__query, got %s", bt3.Name()) } // Server name with hyphens → sanitized to underscores - bt4 := NewBridgeTool("my-server", mcpTool, nil, "", 0, nil) + bt4 := NewBridgeTool("my-server", mcpTool, nil, "", 0, nil, uuid.Nil, nil) if bt4.Name() != "mcp_my_server__query" { t.Errorf("expected name=mcp_my_server__query, got %s", bt4.Name()) } diff --git a/internal/mcp/grant_checker.go b/internal/mcp/grant_checker.go new file mode 100644 index 00000000..e97d571c --- /dev/null +++ b/internal/mcp/grant_checker.go @@ -0,0 +1,160 @@ +package mcp + +import ( + "context" + "log/slog" + "sync" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/bus" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/pkg/protocol" +) + +// 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 +} + +// cacheKey identifies a unique agent+user combination for caching. +type cacheKey struct { + agentID uuid.UUID + userID string +} + +// cacheEntry holds the resolved access info for an agent+user. +type cacheEntry struct { + // allowByServer maps serverID → toolAllowSet. + // nil toolAllowSet means "all tools allowed" (no tool_allow filter). + // non-nil empty map means "no tools allowed" (grant exists but tool_allow is empty). + // Absent key means "server not accessible" (no grant). + allowByServer map[uuid.UUID]map[string]bool +} + +// storeGrantChecker implements GrantChecker backed by MCPServerStore. +// Uses sync.Map for caching with event-bus invalidation (no TTL). +type storeGrantChecker struct { + store store.MCPServerStore + cache sync.Map // map[cacheKey]*cacheEntry +} + +// NewStoreGrantChecker creates a GrantChecker backed by the MCP store. +// Subscribes to TopicCacheMCP for cache invalidation on grant changes. +func NewStoreGrantChecker(mcpStore store.MCPServerStore, msgBus *bus.MessageBus) *storeGrantChecker { + gc := &storeGrantChecker{ + store: mcpStore, + } + + // Subscribe to cache invalidation events. + // When MCP grants are revoked/modified, the HTTP handler broadcasts to this topic. + if msgBus != nil { + msgBus.Subscribe(bus.TopicCacheMCP, func(event bus.Event) { + if event.Name == protocol.EventCacheInvalidate { + gc.cache.Clear() + slog.Debug("mcp.grant_checker.cache_cleared", "trigger", "bus_event") + } + }) + } + + return gc +} + +// 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 { + if gc.store == nil { + // No store = config-path mode, skip grant check + return true + } + + key := cacheKey{agentID: agentID, userID: userID} + + // Try cache first + if cached, ok := gc.cache.Load(key); ok { + entry := cached.(*cacheEntry) + return gc.checkAccess(entry, serverID, toolName, agentID, userID) + } + + // Cache miss — query store + entry, err := gc.loadEntry(ctx, agentID, userID) + if err != nil { + slog.Warn("mcp.grant_checker.load_failed", "agent", agentID, "user", userID, "error", err) + // On error, deny access (fail-closed) + return false + } + + gc.cache.Store(key, entry) + return gc.checkAccess(entry, serverID, toolName, agentID, userID) +} + +// loadEntry queries the store for accessible servers and builds a cache entry. +func (gc *storeGrantChecker) loadEntry(ctx context.Context, agentID uuid.UUID, userID string) (*cacheEntry, error) { + accessible, err := gc.store.ListAccessible(ctx, agentID, userID) + if err != nil { + return nil, err + } + + entry := &cacheEntry{ + allowByServer: make(map[uuid.UUID]map[string]bool, len(accessible)), + } + + for _, info := range accessible { + serverID := info.Server.ID + + // Build tool allow set from grant's tool_allow field + if len(info.ToolAllow) == 0 { + // No tool_allow filter = all tools allowed + entry.allowByServer[serverID] = nil + } else { + // Build set from tool_allow list + allowSet := make(map[string]bool, len(info.ToolAllow)) + for _, t := range info.ToolAllow { + allowSet[t] = true + } + entry.allowByServer[serverID] = allowSet + } + } + + return entry, nil +} + +// 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 { + allowSet, hasServer := entry.allowByServer[serverID] + if !hasServer { + slog.Warn("security.mcp_grant_revoked_at_execute", + "agent", agentID, + "user", userID, + "server", serverID, + "tool", toolName, + "reason", "server_not_accessible", + ) + return false + } + + // nil allowSet means "all tools allowed" + if allowSet == nil { + return true + } + + // Check if tool is in the allow set + if !allowSet[toolName] { + slog.Warn("security.mcp_grant_revoked_at_execute", + "agent", agentID, + "user", userID, + "server", serverID, + "tool", toolName, + "reason", "tool_not_in_allow_list", + ) + return false + } + + return true +} + +// Invalidate clears the entire cache. Called when grants are modified. +func (gc *storeGrantChecker) Invalidate() { + gc.cache.Clear() +} diff --git a/internal/mcp/grant_checker_test.go b/internal/mcp/grant_checker_test.go new file mode 100644 index 00000000..f6ff5a3e --- /dev/null +++ b/internal/mcp/grant_checker_test.go @@ -0,0 +1,189 @@ +package mcp + +import ( + "context" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +// mockMCPStore implements store.MCPServerStore for testing. +type mockMCPStore struct { + accessible []store.MCPAccessInfo + callCount int32 +} + +func (m *mockMCPStore) ListAccessible(ctx context.Context, agentID uuid.UUID, userID string) ([]store.MCPAccessInfo, error) { + atomic.AddInt32(&m.callCount, 1) + return m.accessible, nil +} + +// Stub implementations for interface compliance (not used in grant_checker tests) +func (m *mockMCPStore) CreateServer(ctx context.Context, s *store.MCPServerData) error { return nil } +func (m *mockMCPStore) GetServer(ctx context.Context, id uuid.UUID) (*store.MCPServerData, error) { + return nil, nil +} +func (m *mockMCPStore) GetServerByName(ctx context.Context, name string) (*store.MCPServerData, error) { + return nil, nil +} +func (m *mockMCPStore) ListServers(ctx context.Context) ([]store.MCPServerData, error) { return nil, nil } +func (m *mockMCPStore) UpdateServer(ctx context.Context, id uuid.UUID, updates map[string]any) error { + return nil +} +func (m *mockMCPStore) DeleteServer(ctx context.Context, id uuid.UUID) error { return nil } +func (m *mockMCPStore) GrantToAgent(ctx context.Context, g *store.MCPAgentGrant) error { return nil } +func (m *mockMCPStore) RevokeFromAgent(ctx context.Context, serverID, agentID uuid.UUID) error { + return nil +} +func (m *mockMCPStore) ListAgentGrants(ctx context.Context, agentID uuid.UUID) ([]store.MCPAgentGrant, error) { + return nil, nil +} +func (m *mockMCPStore) ListServerGrants(ctx context.Context, serverID uuid.UUID) ([]store.MCPAgentGrant, error) { + return nil, nil +} +func (m *mockMCPStore) GrantToUser(ctx context.Context, g *store.MCPUserGrant) error { return nil } +func (m *mockMCPStore) RevokeFromUser(ctx context.Context, serverID uuid.UUID, userID string) error { + return nil +} +func (m *mockMCPStore) CountAgentGrantsByServer(ctx context.Context) (map[uuid.UUID]int, error) { + return nil, nil +} +func (m *mockMCPStore) CreateRequest(ctx context.Context, req *store.MCPAccessRequest) error { + return nil +} +func (m *mockMCPStore) ListPendingRequests(ctx context.Context) ([]store.MCPAccessRequest, error) { + return nil, nil +} +func (m *mockMCPStore) ReviewRequest(ctx context.Context, requestID uuid.UUID, approved bool, reviewedBy, note string) error { + return nil +} +func (m *mockMCPStore) GetUserCredentials(ctx context.Context, serverID uuid.UUID, userID string) (*store.MCPUserCredentials, error) { + return nil, nil +} +func (m *mockMCPStore) SetUserCredentials(ctx context.Context, serverID uuid.UUID, userID string, creds store.MCPUserCredentials) error { + return nil +} +func (m *mockMCPStore) DeleteUserCredentials(ctx context.Context, serverID uuid.UUID, userID string) error { + return nil +} + +func TestStoreGrantChecker_CacheHit(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"}}, + }, + } + + gc := NewStoreGrantChecker(mockStore, nil) + ctx := context.Background() + + // First call — cache miss, queries store + result1 := gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool") + if !result1 { + t.Error("expected allowed for accessible server") + } + 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") + if !result2 { + t.Error("expected allowed for accessible server (cached)") + } + if mockStore.callCount != 1 { + t.Errorf("expected 1 DB call (cache hit), got %d", mockStore.callCount) + } +} + +func TestStoreGrantChecker_InvalidateClearsCache(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"}}, + }, + } + + gc := NewStoreGrantChecker(mockStore, nil) + ctx := context.Background() + + // Warm the cache + gc.IsAllowed(ctx, agentID, userID, serverID, "tool") + if mockStore.callCount != 1 { + t.Fatalf("expected 1 call after first IsAllowed, got %d", mockStore.callCount) + } + + // Invalidate cache + gc.Invalidate() + + // Next call should query store again + gc.IsAllowed(ctx, agentID, userID, serverID, "tool") + if mockStore.callCount != 2 { + t.Errorf("expected 2 calls after invalidation, got %d", mockStore.callCount) + } +} + +func TestStoreGrantChecker_ToolAllowFilter(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{"allowed_tool", "another_allowed"}, + }, + }, + } + + gc := NewStoreGrantChecker(mockStore, nil) + ctx := context.Background() + + // Allowed tool + if !gc.IsAllowed(ctx, agentID, userID, serverID, "allowed_tool") { + t.Error("expected true for tool in allow list") + } + + // Not allowed tool + if gc.IsAllowed(ctx, agentID, userID, serverID, "blocked_tool") { + t.Error("expected false for tool not in allow list") + } +} + +func TestStoreGrantChecker_NoGrant_Denied(t *testing.T) { + serverID := uuid.New() + agentID := uuid.New() + userID := "test-user" + + mockStore := &mockMCPStore{ + accessible: []store.MCPAccessInfo{}, // No accessible servers + } + + gc := NewStoreGrantChecker(mockStore, nil) + ctx := context.Background() + + // No grant = denied + if gc.IsAllowed(ctx, agentID, userID, serverID, "any_tool") { + t.Error("expected false when no grant exists") + } +} + +func TestStoreGrantChecker_NilStore_AllowsAll(t *testing.T) { + gc := NewStoreGrantChecker(nil, nil) + ctx := context.Background() + + // Nil store = config-path mode, skip check + if !gc.IsAllowed(ctx, uuid.New(), "user", uuid.New(), "any_tool") { + t.Error("expected true when store is nil (config-path mode)") + } +} diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index 903e693a..33c9c0dc 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -99,6 +99,9 @@ type Manager struct { // DB-backed servers store store.MCPServerStore + // Grant checker for runtime grant verification (nil = skip check) + grantChecker GrantChecker + // Shared connection pool (nil = config-only mode) pool *Pool poolServers map[string]struct{} // server names acquired from pool (for cleanup) @@ -141,6 +144,14 @@ func WithPool(p *Pool) ManagerOption { } } +// WithGrantChecker sets the grant checker for runtime grant verification. +// When set, BridgeTool.Execute rechecks grants before executing tools. +func WithGrantChecker(gc GrantChecker) ManagerOption { + return func(m *Manager) { + m.grantChecker = gc + } +} + // NewManager creates a new MCP Manager. func NewManager(registry *tools.Registry, opts ...ManagerOption) *Manager { m := &Manager{ @@ -167,7 +178,8 @@ func (m *Manager) Start(ctx context.Context) error { continue } - if err := m.connectServer(ctx, name, cfg.Transport, cfg.Command, cfg.Args, cfg.Env, cfg.URL, resolveEnvVars(cfg.Headers), cfg.ToolPrefix, cfg.TimeoutSec); err != nil { + // Config-path servers have no DB ID — pass uuid.Nil + if err := m.connectServer(ctx, name, cfg.Transport, cfg.Command, cfg.Args, cfg.Env, cfg.URL, resolveEnvVars(cfg.Headers), cfg.ToolPrefix, cfg.TimeoutSec, uuid.Nil); err != nil { slog.Warn("mcp.server.connect_failed", "server", name, "error", err) errs = append(errs, fmt.Sprintf("%s: %v", name, err)) } @@ -273,14 +285,14 @@ func (m *Manager) connectAndFilter(ctx context.Context, rs *resolvedServer) erro // 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); err != nil { + rs.args, rs.env, srv.URL, rs.headers, srv.ToolPrefix, srv.TimeoutSec, srv.ID); 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); err != nil { + srv.ToolPrefix, srv.TimeoutSec, srv.ID); err != nil { return err } } @@ -387,7 +399,7 @@ func (m *Manager) maybeEnterSearchMode() { // Update "mcp" group to only the kept inline names. inlineNames := allNames[:mcpToolInlineMaxCount] - tools.RegisterToolGroup("mcp", inlineNames) + m.registry.RegisterToolGroup("mcp", inlineNames) m.searchMode = true slog.Info("mcp.search_mode.enabled", @@ -459,7 +471,7 @@ func (m *Manager) ActivateTools(names []string) { } m.mu.Unlock() - tools.RegisterToolGroup("mcp", activeNames) + m.registry.RegisterToolGroup("mcp", activeNames) slog.Info("mcp.tools.activated", "tools", activated) } @@ -490,7 +502,7 @@ func (m *Manager) ActivateToolIfDeferred(name string) bool { // Register in registry outside lock (registry has its own sync). m.registry.Register(bt) - tools.RegisterToolGroup("mcp", activeNames) + m.registry.RegisterToolGroup("mcp", activeNames) slog.Info("mcp.tools.activated", "tools", []string{name}) return true } diff --git a/internal/mcp/manager_connect.go b/internal/mcp/manager_connect.go index 03a60257..4613a0b0 100644 --- a/internal/mcp/manager_connect.go +++ b/internal/mcp/manager_connect.go @@ -11,7 +11,6 @@ import ( mcpclient "github.com/mark3labs/mcp-go/client" "github.com/mark3labs/mcp-go/client/transport" mcpgo "github.com/mark3labs/mcp-go/mcp" - "github.com/nextlevelbuilder/goclaw/internal/tools" ) // connectAndDiscover creates a client, initializes the MCP handshake, and @@ -101,14 +100,15 @@ func connectAndDiscover(ctx context.Context, name, transportType, command string } // connectServer creates a client, initializes the connection, discovers tools, and registers them. -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) error { +// serverID is the MCP server UUID from DB (uuid.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) 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) + registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec, serverID) ss.toolNames = registeredNames // Create health monitoring context @@ -121,7 +121,7 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman m.mu.Unlock() if len(registeredNames) > 0 { - tools.RegisterToolGroup("mcp:"+name, registeredNames) + m.registry.RegisterToolGroup("mcp:"+name, registeredNames) m.updateMCPGroup() } @@ -138,10 +138,11 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman // registerBridgeTools creates BridgeTools from MCP tool definitions and // registers them in the Manager's registry. Returns registered tool names. -func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int) []string { +// serverID is the MCP server UUID (uuid.Nil for config-path servers). +func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID) []string { var registeredNames []string for _, mcpTool := range mcpTools { - bt := NewBridgeTool(serverName, mcpTool, &ss.clientPtr, toolPrefix, timeoutSec, &ss.connected) + bt := NewBridgeTool(serverName, mcpTool, &ss.clientPtr, toolPrefix, timeoutSec, &ss.connected, serverID, m.grantChecker) if _, exists := m.registry.Get(bt.Name()); exists { slog.Warn("mcp.tool.name_collision", @@ -160,14 +161,15 @@ func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, se // connectViaPool acquires a shared connection from the pool and creates // per-agent BridgeTools pointing to the shared client/connected pointers. -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) error { +// serverID is the MCP server UUID from DB. +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) 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) + registeredNames := m.registerPoolBridgeTools(entry, name, toolPrefix, timeoutSec, serverID) // Track server state and per-agent tool names. // poolServers/poolToolNames keyed by plain name for Close() iteration. @@ -189,7 +191,7 @@ func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name, m.mu.Unlock() if len(registeredNames) > 0 { - tools.RegisterToolGroup("mcp:"+name, registeredNames) + m.registry.RegisterToolGroup("mcp:"+name, registeredNames) m.updateMCPGroup() } @@ -204,10 +206,11 @@ func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name, // registerPoolBridgeTools creates BridgeTools from pool entry's discovered tools, // pointing to the shared client/connected pointers. Returns registered tool names. -func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int) []string { +// serverID is the MCP server UUID from DB. +func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int, serverID uuid.UUID) []string { var registeredNames []string for _, mcpTool := range entry.tools { - bt := NewBridgeTool(serverName, mcpTool, &entry.state.clientPtr, toolPrefix, timeoutSec, &entry.state.connected) + bt := NewBridgeTool(serverName, mcpTool, &entry.state.clientPtr, toolPrefix, timeoutSec, &entry.state.connected, serverID, m.grantChecker) if _, exists := m.registry.Get(bt.Name()); exists { slog.Warn("mcp.tool.name_collision", diff --git a/internal/mcp/manager_lazy_test.go b/internal/mcp/manager_lazy_test.go index ed00a125..63c74cc0 100644 --- a/internal/mcp/manager_lazy_test.go +++ b/internal/mcp/manager_lazy_test.go @@ -5,6 +5,7 @@ import ( "sync" "testing" + "github.com/google/uuid" mcpgo "github.com/mark3labs/mcp-go/mcp" "github.com/nextlevelbuilder/goclaw/internal/tools" ) @@ -15,7 +16,7 @@ func makeBridgeTool(serverName, toolName string) *BridgeTool { Name: toolName, Description: "test tool " + toolName, InputSchema: mcpgo.ToolInputSchema{Type: "object"}, - }, nil, "", 30, nil) + }, nil, "", 30, nil, uuid.Nil, nil) } // setupSearchModeManager returns a Manager already in search mode with the given diff --git a/internal/mcp/manager_tools.go b/internal/mcp/manager_tools.go index e4f7561b..d96773d0 100644 --- a/internal/mcp/manager_tools.go +++ b/internal/mcp/manager_tools.go @@ -8,7 +8,6 @@ import ( mcpgo "github.com/mark3labs/mcp-go/mcp" "github.com/nextlevelbuilder/goclaw/internal/store" - "github.com/nextlevelbuilder/goclaw/internal/tools" ) // UserCredServers returns servers requiring per-user credentials. @@ -53,9 +52,9 @@ func (m *Manager) ServerToolNames(serverName string) []string { func (m *Manager) updateMCPGroup() { allNames := m.ToolNames() if len(allNames) > 0 { - tools.RegisterToolGroup("mcp", allNames) + m.registry.RegisterToolGroup("mcp", allNames) } else { - tools.UnregisterToolGroup("mcp") + m.registry.UnregisterToolGroup("mcp") } } @@ -88,7 +87,7 @@ func (m *Manager) unregisterAllTools() { m.registry.Unregister(toolName) } } - tools.UnregisterToolGroup("mcp:" + name) + m.registry.UnregisterToolGroup("mcp:" + name) slog.Debug("mcp.server.unregistered", "server", name) } @@ -105,7 +104,7 @@ func (m *Manager) unregisterAllTools() { m.servers = make(map[string]*serverState) m.poolServers = nil m.poolToolNames = nil - tools.UnregisterToolGroup("mcp") + m.registry.UnregisterToolGroup("mcp") } // ToolInfo holds a tool's name and description for API responses. diff --git a/internal/tools/policy.go b/internal/tools/policy.go index 0e313959..3e8195b3 100644 --- a/internal/tools/policy.go +++ b/internal/tools/policy.go @@ -10,11 +10,9 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/providers" ) -// toolGroupsMu protects toolGroups from concurrent access. -var toolGroupsMu sync.RWMutex - -// Tool groups map group names to tool names. -var toolGroups = map[string][]string{ +// builtinToolGroups is const-like seed data for per-Registry tool groups. +// Do NOT modify at runtime — each Registry gets a deep copy in NewRegistry(). +var builtinToolGroups = map[string][]string{ "memory": {"memory_search", "memory_get"}, "web": {"web_search", "web_fetch"}, "fs": {"read_file", "write_file", "list_files", "edit"}, @@ -42,39 +40,8 @@ var toolGroups = map[string][]string{ }, } -// RegisterToolGroup adds or replaces a dynamic tool group. -// Used by the MCP manager to register "mcp" and "mcp:{serverName}" groups. -func RegisterToolGroup(name string, members []string) { - toolGroupsMu.Lock() - toolGroups[name] = members - toolGroupsMu.Unlock() -} - -// MergeToolGroup adds members to an existing tool group without removing existing entries. -// Used by per-user MCP tool loading to extend the "mcp" group incrementally. -func MergeToolGroup(name string, members []string) { - toolGroupsMu.Lock() - defer toolGroupsMu.Unlock() - existing := toolGroups[name] - seen := make(map[string]bool, len(existing)) - for _, m := range existing { - seen[m] = true - } - for _, m := range members { - if !seen[m] { - existing = append(existing, m) - seen[m] = true - } - } - toolGroups[name] = existing -} - -// UnregisterToolGroup removes a dynamic tool group. -func UnregisterToolGroup(name string) { - toolGroupsMu.Lock() - delete(toolGroups, name) - toolGroupsMu.Unlock() -} +// Package-level wrappers are REMOVED — use Registry methods instead. +// See Registry.RegisterToolGroup, Registry.MergeToolGroup, Registry.UnregisterToolGroup. // Tool profiles define preset allow sets. var toolProfiles = map[string][]string{ @@ -225,6 +192,11 @@ func (pe *PolicyEngine) evaluate( ) []string { g := pe.globalPolicy + // Get registry for group expansion (may be nil in early boot) + pe.mu.RLock() + reg := pe.registry + pe.mu.RUnlock() + // Step 1: Global profile allowed := pe.applyProfile(allTools, g.Profile) @@ -237,49 +209,49 @@ func (pe *PolicyEngine) evaluate( // Step 3: Global allow list (restricts to only these) if len(g.Allow) > 0 { - allowed = intersectWithSpec(allowed, g.Allow) + allowed = intersectWithSpec(reg, allowed, g.Allow) } // Step 4: Provider-level allow override if g.ByProvider != nil { if pp, ok := g.ByProvider[providerName]; ok && len(pp.Allow) > 0 { - allowed = intersectWithSpec(allowed, pp.Allow) + allowed = intersectWithSpec(reg, allowed, pp.Allow) } } // Step 5: Per-agent allow if agentToolPolicy != nil && len(agentToolPolicy.Allow) > 0 { - allowed = intersectWithSpec(allowed, agentToolPolicy.Allow) + allowed = intersectWithSpec(reg, allowed, agentToolPolicy.Allow) } // Step 6: Per-agent per-provider allow if agentToolPolicy != nil && agentToolPolicy.ByProvider != nil { if pp, ok := agentToolPolicy.ByProvider[providerName]; ok && len(pp.Allow) > 0 { - allowed = intersectWithSpec(allowed, pp.Allow) + allowed = intersectWithSpec(reg, allowed, pp.Allow) } } // Step 7: Group-level allow if len(groupToolAllow) > 0 { - allowed = intersectWithSpec(allowed, groupToolAllow) + allowed = intersectWithSpec(reg, allowed, groupToolAllow) } // Apply global deny if len(g.Deny) > 0 { - allowed = subtractSpec(allowed, g.Deny) + allowed = subtractSpec(reg, allowed, g.Deny) } // Apply agent deny if agentToolPolicy != nil && len(agentToolPolicy.Deny) > 0 { - allowed = subtractSpec(allowed, agentToolPolicy.Deny) + allowed = subtractSpec(reg, allowed, agentToolPolicy.Deny) } // Apply alsoAllow (additive — adds back tools without removing existing) if len(g.AlsoAllow) > 0 { - allowed = unionWithSpec(allowed, allTools, g.AlsoAllow) + allowed = unionWithSpec(reg, allowed, allTools, g.AlsoAllow) } if agentToolPolicy != nil && len(agentToolPolicy.AlsoAllow) > 0 { - allowed = unionWithSpec(allowed, allTools, agentToolPolicy.AlsoAllow) + allowed = unionWithSpec(reg, allowed, allTools, agentToolPolicy.AlsoAllow) } return allowed @@ -298,31 +270,34 @@ func (pe *PolicyEngine) applyProfile(allTools []string, profile string) []string return copySlice(allTools) } - return expandSpec(allTools, spec) + // Get registry for group expansion + pe.mu.RLock() + reg := pe.registry + pe.mu.RUnlock() + + return expandSpec(reg, allTools, spec) } -// --- Set operations with group expansion --- +// --- Set operations with group expansion (using per-Registry tool groups) --- // expandSpec expands a spec list (which may contain "group:xxx") into concrete tool names, -// filtered against available tools. -func expandSpec(available []string, spec []string) []string { - toolGroupsMu.RLock() - defer toolGroupsMu.RUnlock() +// filtered against available tools. Uses per-Registry tool groups to avoid cross-agent races. +func expandSpec(reg *Registry, available []string, spec []string) []string { + if reg == nil { + // Fallback: no group expansion + return expandSpecNoGroups(available, spec) + } + return reg.ExpandToolGroups(available, spec) +} +// expandSpecNoGroups is a fallback when no registry is available. +func expandSpecNoGroups(available []string, spec []string) []string { expanded := make(map[string]bool) for _, s := range spec { - if after, ok := strings.CutPrefix(s, "group:"); ok { - groupName := after - if members, ok := toolGroups[groupName]; ok { - for _, m := range members { - expanded[m] = true - } - } - } else { + if !strings.HasPrefix(s, "group:") { expanded[s] = true } } - var result []string for _, t := range available { if expanded[t] { @@ -333,15 +308,17 @@ func expandSpec(available []string, spec []string) []string { } // intersectWithSpec keeps only tools in `current` that match the spec (with group expansion). -func intersectWithSpec(current []string, spec []string) []string { - toolGroupsMu.RLock() - defer toolGroupsMu.RUnlock() +func intersectWithSpec(reg *Registry, current []string, spec []string) []string { + if reg == nil { + return intersectWithSpecNoGroups(current, spec) + } + reg.toolGroupsMu.RLock() + defer reg.toolGroupsMu.RUnlock() expanded := make(map[string]bool) for _, s := range spec { if after, ok := strings.CutPrefix(s, "group:"); ok { - groupName := after - if members, ok := toolGroups[groupName]; ok { + if members, ok := reg.toolGroups[after]; ok { for _, m := range members { expanded[m] = true } @@ -360,16 +337,34 @@ func intersectWithSpec(current []string, spec []string) []string { return result } +func intersectWithSpecNoGroups(current []string, spec []string) []string { + expanded := make(map[string]bool) + for _, s := range spec { + if !strings.HasPrefix(s, "group:") { + expanded[s] = true + } + } + var result []string + for _, t := range current { + if expanded[t] { + result = append(result, t) + } + } + return result +} + // subtractSpec removes tools matching the spec (with group expansion) from current. -func subtractSpec(current []string, spec []string) []string { - toolGroupsMu.RLock() - defer toolGroupsMu.RUnlock() +func subtractSpec(reg *Registry, current []string, spec []string) []string { + if reg == nil { + return subtractSpecNoGroups(current, spec) + } + reg.toolGroupsMu.RLock() + defer reg.toolGroupsMu.RUnlock() denied := make(map[string]bool) for _, s := range spec { if after, ok := strings.CutPrefix(s, "group:"); ok { - groupName := after - if members, ok := toolGroups[groupName]; ok { + if members, ok := reg.toolGroups[after]; ok { for _, m := range members { denied[m] = true } @@ -388,6 +383,22 @@ func subtractSpec(current []string, spec []string) []string { return result } +func subtractSpecNoGroups(current []string, spec []string) []string { + denied := make(map[string]bool) + for _, s := range spec { + if !strings.HasPrefix(s, "group:") { + denied[s] = true + } + } + var result []string + for _, t := range current { + if !denied[t] { + result = append(result, t) + } + } + return result +} + // subtractSet removes exact tool names from current. func subtractSet(current []string, deny []string) []string { denied := make(map[string]bool, len(deny)) @@ -404,13 +415,13 @@ func subtractSet(current []string, deny []string) []string { } // unionWithSpec adds tools matching spec (from allTools) to current set. -func unionWithSpec(current []string, allTools []string, spec []string) []string { +func unionWithSpec(reg *Registry, current []string, allTools []string, spec []string) []string { existing := make(map[string]bool, len(current)) for _, t := range current { existing[t] = true } - toAdd := expandSpec(allTools, spec) + toAdd := expandSpec(reg, allTools, spec) for _, t := range toAdd { if !existing[t] { current = append(current, t) @@ -423,13 +434,17 @@ func unionWithSpec(current []string, allTools []string, spec []string) []string // IsDenied checks if a tool name is explicitly denied by global or agent policy. // Used to prevent lazy-activated deferred tools from bypassing the deny list. func (pe *PolicyEngine) IsDenied(name string, agentPolicy *config.ToolPolicySpec) bool { + pe.mu.RLock() + reg := pe.registry + pe.mu.RUnlock() + if pe.globalPolicy != nil { - if matchDenySpec(name, pe.globalPolicy.Deny) { + if matchDenySpec(reg, name, pe.globalPolicy.Deny) { return true } } if agentPolicy != nil { - if matchDenySpec(name, agentPolicy.Deny) { + if matchDenySpec(reg, name, agentPolicy.Deny) { return true } } @@ -437,22 +452,12 @@ func (pe *PolicyEngine) IsDenied(name string, agentPolicy *config.ToolPolicySpec } // matchDenySpec returns true if name matches any entry in the deny spec (with group expansion). -func matchDenySpec(name string, spec []string) bool { - toolGroupsMu.RLock() - defer toolGroupsMu.RUnlock() - - for _, s := range spec { - if after, ok := strings.CutPrefix(s, "group:"); ok { - if members, ok := toolGroups[after]; ok { - if slices.Contains(members, name) { - return true - } - } - } else if s == name { - return true - } +func matchDenySpec(reg *Registry, name string, spec []string) bool { + if reg == nil { + // No groups to expand — plain match only + return slices.Contains(spec, name) } - return false + return reg.MatchDenySpec(name, spec) } // StripToolPrefix removes a prefix pattern from a tool name returned by the LLM. diff --git a/internal/tools/policy_mcp_test.go b/internal/tools/policy_mcp_test.go index 8e710f31..64053757 100644 --- a/internal/tools/policy_mcp_test.go +++ b/internal/tools/policy_mcp_test.go @@ -2,11 +2,13 @@ package tools import "testing" -func TestRegisterToolGroup(t *testing.T) { - // Register a new MCP group - RegisterToolGroup("mcp:postgres", []string{"mcp_pg__query", "mcp_pg__list_tables"}) +func TestRegistry_RegisterToolGroup(t *testing.T) { + reg := NewRegistry() - members, ok := toolGroups["mcp:postgres"] + // Register a new MCP group + reg.RegisterToolGroup("mcp:postgres", []string{"mcp_pg__query", "mcp_pg__list_tables"}) + + members, ok := reg.GetToolGroup("mcp:postgres") if !ok { t.Fatal("expected mcp:postgres group to exist") } @@ -15,25 +17,26 @@ func TestRegisterToolGroup(t *testing.T) { } // Unregister - UnregisterToolGroup("mcp:postgres") - if _, ok := toolGroups["mcp:postgres"]; ok { + reg.UnregisterToolGroup("mcp:postgres") + if _, ok := reg.GetToolGroup("mcp:postgres"); ok { t.Error("expected mcp:postgres group to be removed") } } -func TestRegisterToolGroup_UsedInExpand(t *testing.T) { - RegisterToolGroup("mcp:test", []string{"mcp_test__tool_a", "mcp_test__tool_b"}) - defer UnregisterToolGroup("mcp:test") +func TestRegistry_RegisterToolGroup_UsedInExpand(t *testing.T) { + reg := NewRegistry() + reg.RegisterToolGroup("mcp:test", []string{"mcp_test__tool_a", "mcp_test__tool_b"}) + defer reg.UnregisterToolGroup("mcp:test") available := []string{"mcp_test__tool_a", "mcp_test__tool_b", "read_file", "exec"} - expanded := expandSpec(available, []string{"group:mcp:test"}) + expanded := expandSpec(reg, available, []string{"group:mcp:test"}) if len(expanded) != 2 { t.Errorf("expected 2 tools from group:mcp:test, got %d: %v", len(expanded), expanded) } // Verify it works with subtractSpec too - remaining := subtractSpec(available, []string{"group:mcp:test"}) + remaining := subtractSpec(reg, available, []string{"group:mcp:test"}) if len(remaining) != 2 { t.Errorf("expected 2 remaining after subtract, got %d: %v", len(remaining), remaining) } diff --git a/internal/tools/policy_race_test.go b/internal/tools/policy_race_test.go new file mode 100644 index 00000000..23c725ab --- /dev/null +++ b/internal/tools/policy_race_test.go @@ -0,0 +1,171 @@ +package tools + +import ( + "sync" + "testing" +) + +// TestToolGroups_PerRegistry_Isolation verifies that each Registry has isolated +// tool groups. When two registries (representing two agent Loops) register their +// MCP tools to their own "mcp" group, they don't pollute each other. +// +// This test PASSES after Phase 03 fix because toolGroups is now per-Registry. +func TestToolGroups_PerRegistry_Isolation(t *testing.T) { + // Create two separate registries (simulating two agent Loops) + regA := NewRegistry() + regB := NewRegistry() + + // Agent A registers its MCP tools to its registry + toolsA := []string{"mcp_serverA__toolA1", "mcp_serverA__toolA2"} + regA.RegisterToolGroup("mcp", toolsA) + + // Agent B registers its MCP tools to its registry + toolsB := []string{"mcp_serverB__toolB1", "mcp_serverB__toolB2"} + regB.RegisterToolGroup("mcp", toolsB) + + // Agent A expands group:mcp from their registry — should see ONLY toolsA + allToolsA := append(append([]string{}, toolsA...), toolsB...) + expandedA := regA.ExpandToolGroups(allToolsA, []string{"group:mcp"}) + + // Agent B expands group:mcp from their registry — should see ONLY toolsB + allToolsB := append(append([]string{}, toolsA...), toolsB...) + expandedB := regB.ExpandToolGroups(allToolsB, []string{"group:mcp"}) + + // Verify isolation: A sees only A's tools + if !containsTool(expandedA, "mcp_serverA__toolA1") || !containsTool(expandedA, "mcp_serverA__toolA2") { + t.Errorf("Agent A should see their own tools, got: %v", expandedA) + } + if containsTool(expandedA, "mcp_serverB__toolB1") || containsTool(expandedA, "mcp_serverB__toolB2") { + t.Errorf("Agent A should NOT see Agent B's tools, got: %v", expandedA) + } + + // Verify isolation: B sees only B's tools + if !containsTool(expandedB, "mcp_serverB__toolB1") || !containsTool(expandedB, "mcp_serverB__toolB2") { + t.Errorf("Agent B should see their own tools, got: %v", expandedB) + } + if containsTool(expandedB, "mcp_serverA__toolA1") || containsTool(expandedB, "mcp_serverA__toolA2") { + t.Errorf("Agent B should NOT see Agent A's tools, got: %v", expandedB) + } +} + +// TestToolGroups_PerRegistry_ConcurrentWrite verifies concurrent writes to +// the same Registry are safe (no data race) but last-writer-wins within +// that registry is expected behavior. +func TestToolGroups_PerRegistry_ConcurrentWrite(t *testing.T) { + reg := NewRegistry() + + var wg sync.WaitGroup + wg.Add(2) + + toolsA := []string{"toolA1", "toolA2"} + toolsB := []string{"toolB1", "toolB2"} + + barrier := make(chan struct{}) + + go func() { + defer wg.Done() + <-barrier + reg.RegisterToolGroup("mcp", toolsA) + }() + + go func() { + defer wg.Done() + <-barrier + reg.RegisterToolGroup("mcp", toolsB) + }() + + close(barrier) + wg.Wait() + + // One of them wins — that's expected within a single registry + // The important thing is: no data race (run with -race flag) + members, ok := reg.GetToolGroup("mcp") + if !ok { + t.Fatal("expected mcp group to exist") + } + if len(members) != 2 { + t.Errorf("expected 2 members (from winner), got %d: %v", len(members), members) + } +} + +// TestToolGroups_Clone_DeepCopy verifies that Clone() deep-copies toolGroups +// so modifications to the clone don't affect the original. +func TestToolGroups_Clone_DeepCopy(t *testing.T) { + orig := NewRegistry() + orig.RegisterToolGroup("mcp", []string{"tool1", "tool2"}) + + clone := orig.Clone() + + // Modify clone's mcp group + clone.RegisterToolGroup("mcp", []string{"tool3", "tool4"}) + + // Original should still have tool1, tool2 + members, ok := orig.GetToolGroup("mcp") + if !ok { + t.Fatal("original mcp group should exist") + } + if !containsTool(members, "tool1") || !containsTool(members, "tool2") { + t.Errorf("original mcp group modified, got: %v", members) + } + if containsTool(members, "tool3") || containsTool(members, "tool4") { + t.Errorf("original mcp group polluted by clone, got: %v", members) + } +} + +// TestToolGroups_MergeToolGroup_Additive verifies MergeToolGroup behavior. +func TestToolGroups_MergeToolGroup_Additive(t *testing.T) { + reg := NewRegistry() + + // First merge + reg.MergeToolGroup("test_merge", []string{"tool1", "tool2"}) + + // Second merge (should add, not replace) + reg.MergeToolGroup("test_merge", []string{"tool2", "tool3"}) + + // Expand + allTools := []string{"tool1", "tool2", "tool3", "tool4"} + expanded := reg.ExpandToolGroups(allTools, []string{"group:test_merge"}) + + // Should have tool1, tool2, tool3 (merged) + if !containsTool(expanded, "tool1") { + t.Error("tool1 missing from merged group") + } + if !containsTool(expanded, "tool2") { + t.Error("tool2 missing from merged group") + } + if !containsTool(expanded, "tool3") { + t.Error("tool3 missing from merged group") + } +} + +// TestToolGroups_BuiltinGroups_Seeded verifies that NewRegistry() seeds +// builtin tool groups from builtinToolGroups. +func TestToolGroups_BuiltinGroups_Seeded(t *testing.T) { + reg := NewRegistry() + + // Check that builtin groups are present + memory, ok := reg.GetToolGroup("memory") + if !ok { + t.Fatal("expected 'memory' builtin group to exist") + } + if !containsTool(memory, "memory_search") { + t.Errorf("memory group should contain memory_search, got: %v", memory) + } + + web, ok := reg.GetToolGroup("web") + if !ok { + t.Fatal("expected 'web' builtin group to exist") + } + if !containsTool(web, "web_search") || !containsTool(web, "web_fetch") { + t.Errorf("web group should contain web_search and web_fetch, got: %v", web) + } +} + +func containsTool(tools []string, name string) bool { + for _, t := range tools { + if t == name { + return true + } + } + return false +} diff --git a/internal/tools/registry.go b/internal/tools/registry.go index a8c77d3a..bc4ee5d3 100644 --- a/internal/tools/registry.go +++ b/internal/tools/registry.go @@ -24,19 +24,30 @@ type Registry struct { rateLimiter *ToolRateLimiter // nil = no rate limiting scrubbing bool // scrub credentials from output (default true) + // Per-registry tool groups (eliminates global map race condition). + // MCP tools register their groups here so each Loop has isolated namespace. + toolGroups map[string][]string + toolGroupsMu sync.RWMutex + // deferredActivator is called when a tool is not in the registry but may be // a deferred MCP tool. Returns true if the tool was successfully activated. deferredActivator func(name string) bool } func NewRegistry() *Registry { - return &Registry{ - tools: make(map[string]Tool), - metadata: make(map[string]ToolMetadata), - aliases: make(map[string]string), - disabled: make(map[string]bool), - scrubbing: true, // enabled by default + r := &Registry{ + tools: make(map[string]Tool), + metadata: make(map[string]ToolMetadata), + aliases: make(map[string]string), + disabled: make(map[string]bool), + toolGroups: make(map[string][]string), + scrubbing: true, // enabled by default } + // Seed built-in tool groups (deep copy from package-level constant data) + for name, members := range builtinToolGroups { + r.toolGroups[name] = append([]string(nil), members...) + } + return r } // SetDeferredActivator registers a callback that activates deferred tools on demand. @@ -337,11 +348,15 @@ func (r *Registry) Count() int { func (r *Registry) Clone() *Registry { r.mu.RLock() defer r.mu.RUnlock() + r.toolGroupsMu.RLock() + defer r.toolGroupsMu.RUnlock() + clone := &Registry{ tools: make(map[string]Tool, len(r.tools)), metadata: make(map[string]ToolMetadata, len(r.metadata)), aliases: make(map[string]string, len(r.aliases)), disabled: make(map[string]bool, len(r.disabled)), + toolGroups: make(map[string][]string, len(r.toolGroups)), rateLimiter: r.rateLimiter, scrubbing: r.scrubbing, } @@ -349,5 +364,102 @@ func (r *Registry) Clone() *Registry { maps.Copy(clone.metadata, r.metadata) maps.Copy(clone.aliases, r.aliases) maps.Copy(clone.disabled, r.disabled) + // Deep-copy toolGroups (each slice must be copied) + for name, members := range r.toolGroups { + clone.toolGroups[name] = append([]string(nil), members...) + } return clone } + +// RegisterToolGroup adds or replaces a dynamic tool group. +// Used by the MCP manager to register "mcp" and "mcp:{serverName}" groups. +func (r *Registry) RegisterToolGroup(name string, members []string) { + r.toolGroupsMu.Lock() + r.toolGroups[name] = members + r.toolGroupsMu.Unlock() +} + +// MergeToolGroup adds members to an existing tool group without removing existing entries. +// Used by per-user MCP tool loading to extend the "mcp" group incrementally. +func (r *Registry) MergeToolGroup(name string, members []string) { + r.toolGroupsMu.Lock() + defer r.toolGroupsMu.Unlock() + existing := r.toolGroups[name] + seen := make(map[string]bool, len(existing)) + for _, m := range existing { + seen[m] = true + } + for _, m := range members { + if !seen[m] { + existing = append(existing, m) + seen[m] = true + } + } + r.toolGroups[name] = existing +} + +// UnregisterToolGroup removes a dynamic tool group. +func (r *Registry) UnregisterToolGroup(name string) { + r.toolGroupsMu.Lock() + delete(r.toolGroups, name) + r.toolGroupsMu.Unlock() +} + +// GetToolGroup returns the members of a tool group (thread-safe). +func (r *Registry) GetToolGroup(name string) ([]string, bool) { + r.toolGroupsMu.RLock() + defer r.toolGroupsMu.RUnlock() + members, ok := r.toolGroups[name] + if !ok { + return nil, false + } + // Return a copy to prevent mutation + return append([]string(nil), members...), true +} + +// ExpandToolGroups expands a spec list (which may contain "group:xxx") into concrete tool names, +// filtered against available tools. Thread-safe. +func (r *Registry) ExpandToolGroups(available []string, spec []string) []string { + r.toolGroupsMu.RLock() + defer r.toolGroupsMu.RUnlock() + + expanded := make(map[string]bool) + for _, s := range spec { + if after, ok := strings.CutPrefix(s, "group:"); ok { + if members, ok := r.toolGroups[after]; ok { + for _, m := range members { + expanded[m] = true + } + } + } else { + expanded[s] = true + } + } + + var result []string + for _, t := range available { + if expanded[t] { + result = append(result, t) + } + } + return result +} + +// MatchDenySpec returns true if name matches any entry in the deny spec (with group expansion). +func (r *Registry) MatchDenySpec(name string, spec []string) bool { + r.toolGroupsMu.RLock() + defer r.toolGroupsMu.RUnlock() + + for _, s := range spec { + if after, ok := strings.CutPrefix(s, "group:"); ok { + if members, ok := r.toolGroups[after]; ok { + if slices.Contains(members, name) { + return true + } + } + } else if s == name { + return true + } + } + return false +} diff --git a/tests/integration/mcp_grant_revoke_test.go b/tests/integration/mcp_grant_revoke_test.go new file mode 100644 index 00000000..8f820dcd --- /dev/null +++ b/tests/integration/mcp_grant_revoke_test.go @@ -0,0 +1,261 @@ +//go:build integration + +package integration + +import ( + "context" + "database/sql" + "sync/atomic" + "testing" + + mcpclient "github.com/mark3labs/mcp-go/client" + mcpgo "github.com/mark3labs/mcp-go/mcp" + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/mcp" + "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/store/pg" +) + +// TestBridgeTool_Execute_RevokeAgentGrant_ReturnsError verifies that after revoking +// an agent grant, BridgeTool.Execute returns an error instead of executing the tool. +// +// This test MUST FAIL initially (Phase 01 TDD) because BridgeTool.Execute currently +// only checks `connected` status — it does NOT recheck grants. +func TestBridgeTool_Execute_RevokeAgentGrant_ReturnsError(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + serverID := seedMCPServer(t, db, tenantID) + + // Grant agent access to the MCP server + grantAgentAccess(t, db, tenantID, serverID, agentID) + + // Create MCP store + mcpStore := pg.NewPGMCPServerStore(db, testEncryptionKey) + ctx := store.WithTenantID(context.Background(), tenantID) + ctx = store.WithAgentID(ctx, agentID) + ctx = store.WithUserID(ctx, "test-user") + + // Verify grant is active + accessible, err := mcpStore.ListAccessible(ctx, agentID, "test-user") + if err != nil { + t.Fatalf("ListAccessible: %v", err) + } + if len(accessible) == 0 { + t.Fatal("expected at least 1 accessible server after grant") + } + + // Create a fake MCP client that returns a stub result + fakeClient := &fakeMCPClient{result: &mcpgo.CallToolResult{ + Content: []mcpgo.Content{mcpgo.TextContent{Text: "success"}}, + }} + + // Create BridgeTool with the fake client + clientPtr := &atomic.Pointer[mcpclient.Client]{} + // Note: We need to cast the fake client to the interface type + // This is a workaround since mcp-go client is a struct, not an interface + connected := &atomic.Bool{} + connected.Store(true) + + // Create a grant checker that checks the store + grantChecker := mcp.NewStoreGrantChecker(mcpStore, nil) + + tool := mcp.NewBridgeTool( + "test-server", + mcpgo.Tool{Name: "test_tool", Description: "test"}, + clientPtr, + "mcp_test", + 60, + connected, + serverID, + grantChecker, + ) + + // Execute should work before revoke (will fail due to nil client, but that's expected) + // The key point is: after revoke, it should return "grant revoked" error + + // Now revoke the agent grant + err = mcpStore.RevokeFromAgent(ctx, serverID, agentID) + if err != nil { + t.Fatalf("RevokeFromAgent: %v", err) + } + + // Execute the tool after revoke + // EXPECTED (after Phase 02 fix): should return ErrorResult with "grant revoked" + // ACTUAL (currently): will try to execute and fail with "no active client" or succeed + result := tool.Execute(ctx, map[string]any{"arg": "value"}) + + // This assertion SHOULD PASS after Phase 02, but FAILS now + // because BridgeTool.Execute does NOT recheck grants + if !result.IsError { + t.Error("expected error result after grant revoked, but got success") + } + if result.IsError && !containsGrantRevoked(result.ForLLM) { + t.Errorf("expected 'grant revoked' error, got: %s", result.ForLLM) + } +} + +// TestBridgeTool_Execute_RevokeUserGrant_ReturnsError verifies that after revoking +// a user grant, BridgeTool.Execute returns an error. +// +// This test MUST FAIL initially (Phase 01 TDD). +func TestBridgeTool_Execute_RevokeUserGrant_ReturnsError(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + serverID := seedMCPServer(t, db, tenantID) + userID := "test-user-" + uuid.New().String()[:8] + + // Grant agent access (required for ListAccessible) + grantAgentAccess(t, db, tenantID, serverID, agentID) + + // Grant user access + grantUserAccess(t, db, tenantID, serverID, userID) + + // Create MCP store + mcpStore := pg.NewPGMCPServerStore(db, testEncryptionKey) + ctx := store.WithTenantID(context.Background(), tenantID) + ctx = store.WithAgentID(ctx, agentID) + ctx = store.WithUserID(ctx, userID) + + // Verify both grants are active + accessible, err := mcpStore.ListAccessible(ctx, agentID, userID) + if err != nil { + t.Fatalf("ListAccessible: %v", err) + } + if len(accessible) == 0 { + t.Fatal("expected accessible server after grants") + } + + // Create BridgeTool + clientPtr := &atomic.Pointer[mcpclient.Client]{} + connected := &atomic.Bool{} + connected.Store(true) + + // Create a grant checker that checks the store + grantChecker := mcp.NewStoreGrantChecker(mcpStore, nil) + + tool := mcp.NewBridgeTool( + "test-server", + mcpgo.Tool{Name: "test_tool", Description: "test"}, + clientPtr, + "mcp_test", + 60, + connected, + serverID, + grantChecker, + ) + + // Revoke the USER grant (agent grant still active) + err = mcpStore.RevokeFromUser(ctx, serverID, userID) + if err != nil { + t.Fatalf("RevokeFromUser: %v", err) + } + + // Execute the tool after user revoke + // EXPECTED (after Phase 02 fix): should return "grant revoked" since user lost access + // ACTUAL (currently): does not check user grants at execute time + result := tool.Execute(ctx, map[string]any{"arg": "value"}) + + // This assertion SHOULD PASS after Phase 02, but FAILS now + if !result.IsError { + t.Error("expected error result after user grant revoked") + } + if result.IsError && !containsGrantRevoked(result.ForLLM) { + t.Errorf("expected 'grant revoked' error, got: %s", result.ForLLM) + } +} + +// TestResolver_Rebuild_AfterRevoke_NoToolInPrompt verifies that after revoking a grant, +// the next resolver.Get() returns a Loop without the revoked tool in the prompt. +// +// This test SHOULD PASS even before fixes (regression guard) because the existing +// unregisterAllTools + fresh clone mechanism already handles prompt rebuild. +func TestResolver_Rebuild_AfterRevoke_NoToolInPrompt(t *testing.T) { + db := testDB(t) + tenantID, agentID := seedTenantAgent(t, db) + serverID := seedMCPServer(t, db, tenantID) + + // Grant agent access + grantAgentAccess(t, db, tenantID, serverID, agentID) + + // Create MCP store + mcpStore := pg.NewPGMCPServerStore(db, testEncryptionKey) + ctx := store.WithTenantID(context.Background(), tenantID) + + // Verify grant is active + accessible, err := mcpStore.ListAccessible(ctx, agentID, "test-user") + if err != nil { + t.Fatalf("ListAccessible before revoke: %v", err) + } + if len(accessible) == 0 { + t.Fatal("expected accessible server before revoke") + } + serverName := accessible[0].Server.Name + + // Revoke the grant + err = mcpStore.RevokeFromAgent(ctx, serverID, agentID) + if err != nil { + t.Fatalf("RevokeFromAgent: %v", err) + } + + // Verify no servers accessible after revoke + accessible, err = mcpStore.ListAccessible(ctx, agentID, "test-user") + if err != nil { + t.Fatalf("ListAccessible after revoke: %v", err) + } + if len(accessible) != 0 { + t.Errorf("expected 0 accessible servers after revoke, got %d", len(accessible)) + } + + // This test passes as a regression guard: + // The next LoadForAgent() will query ListAccessible which returns empty, + // so no MCP tools will be registered. The prompt rebuild mechanism works. + t.Logf("Regression guard PASS: server %q no longer accessible after revoke", serverName) +} + +// --- Helpers --- + +func grantAgentAccess(t *testing.T, db *sql.DB, tenantID, serverID, agentID uuid.UUID) { + t.Helper() + _, err := db.Exec( + `INSERT INTO mcp_agent_grants (id, server_id, agent_id, enabled, granted_by, created_at, tenant_id) + VALUES ($1, $2, $3, true, 'test-admin', NOW(), $4) + ON CONFLICT (server_id, agent_id) DO UPDATE SET enabled = true`, + uuid.New(), serverID, agentID, tenantID) + if err != nil { + t.Fatalf("grantAgentAccess: %v", err) + } +} + +func grantUserAccess(t *testing.T, db *sql.DB, tenantID, serverID uuid.UUID, userID string) { + t.Helper() + _, err := db.Exec( + `INSERT INTO mcp_user_grants (id, server_id, user_id, enabled, granted_by, created_at, tenant_id) + VALUES ($1, $2, $3, true, 'test-admin', NOW(), $4) + ON CONFLICT (server_id, user_id) DO UPDATE SET enabled = true`, + uuid.New(), serverID, userID, tenantID) + if err != nil { + t.Fatalf("grantUserAccess: %v", err) + } +} + +func containsGrantRevoked(s string) bool { + return len(s) > 0 && (contains(s, "grant revoked") || contains(s, "grant denied")) +} + +func contains(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} + +// fakeMCPClient is a stub for testing. Since mcpclient.Client is a struct +// and not an interface, we cannot directly mock it. The test relies on +// the clientPtr being nil or the connection being marked as disconnected. +type fakeMCPClient struct { + result *mcpgo.CallToolResult + err error +}