mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-07-14 05:06:15 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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")})
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user