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
This commit is contained in:
viettranx
2026-04-16 17:07:56 +07:00
parent 8618ed94b2
commit 8b8da3a3dd
20 changed files with 1128 additions and 150 deletions
+3
View File
@@ -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,
+5 -3
View File
@@ -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")
+2 -2
View File
@@ -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
+12 -7
View File
@@ -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,
+7
View File
@@ -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,
+14
View File
@@ -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"`
+14
View File
@@ -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")})
+19 -3
View File
@@ -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))
}
+5 -4
View File
@@ -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())
}
+160
View File
@@ -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()
}
+189
View File
@@ -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)")
}
}
+18 -6
View File
@@ -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
}
+14 -11
View File
@@ -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",
+2 -1
View File
@@ -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
+4 -5
View File
@@ -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.
+96 -91
View File
@@ -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.
+14 -11
View File
@@ -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)
}
+171
View File
@@ -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
}
+118 -6
View File
@@ -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
}
+261
View File
@@ -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
}