mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 08:18:19 +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:
1 parent
8618ed94b2
commit
8b8da3a3dd
20 files changed
+1128
-150
No files matched your search
@@ -127,8 +127,10 @@ func wireExtras(
|
|||||||
|
|
||||||
// 5. Shared MCP connection pool (eliminates duplicate connections across agents)
|
// 5. Shared MCP connection pool (eliminates duplicate connections across agents)
|
||||||
var mcpPool *mcpbridge.Pool
|
var mcpPool *mcpbridge.Pool
|
||||||
|
var mcpGrantChecker mcpbridge.GrantChecker
|
||||||
if stores.MCP != nil {
|
if stores.MCP != nil {
|
||||||
mcpPool = mcpbridge.NewPool(mcpbridge.DefaultPoolConfig())
|
mcpPool = mcpbridge.NewPool(mcpbridge.DefaultPoolConfig())
|
||||||
|
mcpGrantChecker = mcpbridge.NewStoreGrantChecker(stores.MCP, msgBus)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 6. Set up agent resolver: lazy-creates Loops from DB
|
// 6. Set up agent resolver: lazy-creates Loops from DB
|
||||||
@@ -221,6 +223,7 @@ func wireExtras(
|
|||||||
BuiltinToolStore: stores.BuiltinTools,
|
BuiltinToolStore: stores.BuiltinTools,
|
||||||
MCPStore: stores.MCP,
|
MCPStore: stores.MCP,
|
||||||
MCPPool: mcpPool,
|
MCPPool: mcpPool,
|
||||||
|
MCPGrantChecker: mcpGrantChecker,
|
||||||
ConfigPermStore: stores.ConfigPermissions,
|
ConfigPermStore: stores.ConfigPermissions,
|
||||||
MediaStore: mediaStore,
|
MediaStore: mediaStore,
|
||||||
ModelPricing: appCfg.Telemetry.ModelPricing,
|
ModelPricing: appCfg.Telemetry.ModelPricing,
|
||||||
|
|||||||
@@ -337,13 +337,15 @@ func TestLoop_LazyMCP_DenyList_GroupDeny(t *testing.T) {
|
|||||||
return true
|
return true
|
||||||
})
|
})
|
||||||
|
|
||||||
// Register a custom group containing the MCP tool.
|
// Register a custom group containing the MCP tool (using per-Registry method).
|
||||||
tools.RegisterToolGroup("mcp_test", []string{"mcp_svc__get_data"})
|
reg.RegisterToolGroup("mcp_test", []string{"mcp_svc__get_data"})
|
||||||
defer tools.UnregisterToolGroup("mcp_test")
|
defer reg.UnregisterToolGroup("mcp_test")
|
||||||
|
|
||||||
pe := tools.NewPolicyEngine(&config.ToolsConfig{
|
pe := tools.NewPolicyEngine(&config.ToolsConfig{
|
||||||
Deny: []string{"group:mcp_test"},
|
Deny: []string{"group:mcp_test"},
|
||||||
})
|
})
|
||||||
|
// PolicyEngine needs registry to expand group:mcp_test
|
||||||
|
pe.SetRegistry(reg)
|
||||||
|
|
||||||
if !pe.IsDenied("mcp_svc__get_data", nil) {
|
if !pe.IsDenied("mcp_svc__get_data", nil) {
|
||||||
t.Error("tool should be denied via group: pattern")
|
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.
|
// shared tool registry so ExecuteWithContext can resolve them by name.
|
||||||
reg, _ := l.tools.(*tools.Registry)
|
reg, _ := l.tools.(*tools.Registry)
|
||||||
for _, mcpTool := range entry.MCPTools() {
|
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.
|
// Register in registry so ExecuteWithContext can find them.
|
||||||
// Skip if already registered (another user loaded this server with same tool names).
|
// Skip if already registered (another user loaded this server with same tool names).
|
||||||
if reg != nil {
|
if reg != nil {
|
||||||
@@ -109,7 +109,7 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool
|
|||||||
for _, t := range userTools {
|
for _, t := range userTools {
|
||||||
names = append(names, t.Name())
|
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))
|
slog.Info("mcp.user_tools_loaded", "user", userID, "tools", len(userTools))
|
||||||
}
|
}
|
||||||
return userTools
|
return userTools
|
||||||
|
|||||||
@@ -112,6 +112,7 @@ type Loop struct {
|
|||||||
domainBus eventbus.DomainEventBus // V3 domain event bus for consolidation pipeline
|
domainBus eventbus.DomainEventBus // V3 domain event bus for consolidation pipeline
|
||||||
sessions store.SessionStore
|
sessions store.SessionStore
|
||||||
tools tools.ToolExecutor
|
tools tools.ToolExecutor
|
||||||
|
registry *tools.Registry // direct registry access for MergeToolGroup (per-Registry tool groups)
|
||||||
toolPolicy *tools.PolicyEngine // optional: filters tools sent to LLM
|
toolPolicy *tools.PolicyEngine // optional: filters tools sent to LLM
|
||||||
agentToolPolicy *config.ToolPolicySpec // per-agent tool policy from DB (nil = no restrictions)
|
agentToolPolicy *config.ToolPolicySpec // per-agent tool policy from DB (nil = no restrictions)
|
||||||
activeRuns atomic.Int32 // number of currently executing runs
|
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)
|
userSetups sync.Map // userID → *userSetup (workspace + seeding state, per Loop instance)
|
||||||
|
|
||||||
// Per-user MCP tools: servers requiring user credentials get connected per-request.
|
// Per-user MCP tools: servers requiring user credentials get connected per-request.
|
||||||
mcpStore store.MCPServerStore // for credential lookup
|
mcpStore store.MCPServerStore // for credential lookup
|
||||||
mcpPool *mcpbridge.Pool // user-keyed connection pool
|
mcpPool *mcpbridge.Pool // user-keyed connection pool
|
||||||
mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
||||||
mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools)
|
mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools)
|
||||||
|
mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip)
|
||||||
|
|
||||||
// Compaction config (memory flush settings)
|
// Compaction config (memory flush settings)
|
||||||
compactionCfg *config.CompactionConfig
|
compactionCfg *config.CompactionConfig
|
||||||
@@ -416,9 +418,10 @@ type LoopConfig struct {
|
|||||||
MemoryStore store.MemoryStore
|
MemoryStore store.MemoryStore
|
||||||
|
|
||||||
// Per-user MCP tools (servers requiring per-user credentials)
|
// Per-user MCP tools (servers requiring per-user credentials)
|
||||||
MCPStore store.MCPServerStore // for credential lookup
|
MCPStore store.MCPServerStore // for credential lookup
|
||||||
MCPPool *mcpbridge.Pool // user-keyed connection pool
|
MCPPool *mcpbridge.Pool // user-keyed connection pool
|
||||||
MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
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)
|
// V3 orchestration mode (resolved by resolver, controls tool visibility)
|
||||||
OrchMode OrchestrationMode
|
OrchMode OrchestrationMode
|
||||||
@@ -491,6 +494,7 @@ func NewLoop(cfg LoopConfig) *Loop {
|
|||||||
hookDispatcher: cfg.HookDispatcher,
|
hookDispatcher: cfg.HookDispatcher,
|
||||||
sessions: cfg.Sessions,
|
sessions: cfg.Sessions,
|
||||||
tools: cfg.Tools,
|
tools: cfg.Tools,
|
||||||
|
registry: cfg.Tools,
|
||||||
toolPolicy: cfg.ToolPolicy,
|
toolPolicy: cfg.ToolPolicy,
|
||||||
agentToolPolicy: cfg.AgentToolPolicy,
|
agentToolPolicy: cfg.AgentToolPolicy,
|
||||||
onEvent: cfg.OnEvent,
|
onEvent: cfg.OnEvent,
|
||||||
@@ -540,6 +544,7 @@ func NewLoop(cfg LoopConfig) *Loop {
|
|||||||
mcpStore: cfg.MCPStore,
|
mcpStore: cfg.MCPStore,
|
||||||
mcpPool: cfg.MCPPool,
|
mcpPool: cfg.MCPPool,
|
||||||
mcpUserCredSrvs: cfg.MCPUserCredSrvs,
|
mcpUserCredSrvs: cfg.MCPUserCredSrvs,
|
||||||
|
mcpGrantChecker: cfg.MCPGrantChecker,
|
||||||
orchMode: cfg.OrchMode,
|
orchMode: cfg.OrchMode,
|
||||||
delegateTargets: cfg.DelegateTargets,
|
delegateTargets: cfg.DelegateTargets,
|
||||||
evolutionMetricsStore: cfg.EvolutionMetricsStore,
|
evolutionMetricsStore: cfg.EvolutionMetricsStore,
|
||||||
|
|||||||
@@ -80,6 +80,9 @@ type ResolverDeps struct {
|
|||||||
// Shared MCP connection pool — eliminates duplicate connections across agents
|
// Shared MCP connection pool — eliminates duplicate connections across agents
|
||||||
MCPPool *mcpbridge.Pool
|
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
|
// Skill access store — for per-agent skill visibility filtering
|
||||||
SkillAccessStore store.SkillAccessStore
|
SkillAccessStore store.SkillAccessStore
|
||||||
|
|
||||||
@@ -304,6 +307,9 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc {
|
|||||||
if deps.MCPPool != nil {
|
if deps.MCPPool != nil {
|
||||||
mcpOpts = append(mcpOpts, mcpbridge.WithPool(deps.MCPPool))
|
mcpOpts = append(mcpOpts, mcpbridge.WithPool(deps.MCPPool))
|
||||||
}
|
}
|
||||||
|
if deps.MCPGrantChecker != nil {
|
||||||
|
mcpOpts = append(mcpOpts, mcpbridge.WithGrantChecker(deps.MCPGrantChecker))
|
||||||
|
}
|
||||||
mcpMgr := mcpbridge.NewManager(toolsReg, mcpOpts...)
|
mcpMgr := mcpbridge.NewManager(toolsReg, mcpOpts...)
|
||||||
if err := mcpMgr.LoadForAgent(ctx, ag.ID, ""); err != nil {
|
if err := mcpMgr.LoadForAgent(ctx, ag.ID, ""); err != nil {
|
||||||
slog.Warn("failed to load MCP servers for agent", "agent", agentKey, "error", err)
|
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,
|
MCPStore: deps.MCPStore,
|
||||||
MCPPool: deps.MCPPool,
|
MCPPool: deps.MCPPool,
|
||||||
MCPUserCredSrvs: mcpUserCredSrvs,
|
MCPUserCredSrvs: mcpUserCredSrvs,
|
||||||
|
MCPGrantChecker: deps.MCPGrantChecker,
|
||||||
OrchMode: orchMode,
|
OrchMode: orchMode,
|
||||||
DelegateTargets: delegateTargets,
|
DelegateTargets: delegateTargets,
|
||||||
EvolutionMetricsStore: evoMetricsStore,
|
EvolutionMetricsStore: evoMetricsStore,
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ import (
|
|||||||
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/i18n"
|
"github.com/nextlevelbuilder/goclaw/internal/i18n"
|
||||||
|
"github.com/nextlevelbuilder/goclaw/internal/skills"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/store/pg"
|
"github.com/nextlevelbuilder/goclaw/internal/store/pg"
|
||||||
)
|
)
|
||||||
@@ -132,6 +133,19 @@ func (h *SkillsHandler) doSkillsImport(ctx context.Context, r io.Reader, userID
|
|||||||
continue
|
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 {
|
var meta struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Slug string `json:"slug"`
|
Slug string `json:"slug"`
|
||||||
|
|||||||
@@ -106,6 +106,20 @@ func (h *SkillsHandler) handleUpload(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
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)
|
name, description, slug, frontmatter := skills.ParseSkillFrontmatter(skillContent)
|
||||||
if name == "" {
|
if name == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": i18n.T(locale, i18n.MsgRequired, "name in SKILL.md frontmatter")})
|
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"
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
mcpclient "github.com/mark3labs/mcp-go/client"
|
mcpclient "github.com/mark3labs/mcp-go/client"
|
||||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||||
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -19,14 +21,16 @@ import (
|
|||||||
// safe reconnection without data races.
|
// safe reconnection without data races.
|
||||||
type BridgeTool struct {
|
type BridgeTool struct {
|
||||||
serverName string
|
serverName string
|
||||||
toolName string // original MCP tool name
|
serverID uuid.UUID // MCP server ID (for grant recheck)
|
||||||
registeredName string // may include prefix: "{prefix}__{toolName}"
|
toolName string // original MCP tool name
|
||||||
|
registeredName string // may include prefix: "{prefix}__{toolName}"
|
||||||
description string
|
description string
|
||||||
inputSchema map[string]any // JSON Schema for parameters
|
inputSchema map[string]any // JSON Schema for parameters
|
||||||
requiredSet map[string]bool
|
requiredSet map[string]bool
|
||||||
clientPtr *atomic.Pointer[mcpclient.Client] // shared with serverState for atomic swap on reconnect
|
clientPtr *atomic.Pointer[mcpclient.Client] // shared with serverState for atomic swap on reconnect
|
||||||
timeoutSec int
|
timeoutSec int
|
||||||
connected *atomic.Bool
|
connected *atomic.Bool
|
||||||
|
grantChecker GrantChecker // for runtime grant recheck (nil = skip check)
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewBridgeTool creates a BridgeTool from an MCP Tool definition.
|
// 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.
|
// If prefix is empty, it is auto-derived from the server name.
|
||||||
// clientPtr is a shared atomic pointer from serverState — reconnection swaps it
|
// clientPtr is a shared atomic pointer from serverState — reconnection swaps it
|
||||||
// atomically, and all BridgeTools see the new client without explicit notification.
|
// 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
|
name := mcpTool.Name
|
||||||
effectivePrefix := ensureMCPPrefix(prefix, serverName)
|
effectivePrefix := ensureMCPPrefix(prefix, serverName)
|
||||||
registered := effectivePrefix + "__" + name
|
registered := effectivePrefix + "__" + name
|
||||||
@@ -52,6 +57,7 @@ func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Poin
|
|||||||
|
|
||||||
return &BridgeTool{
|
return &BridgeTool{
|
||||||
serverName: serverName,
|
serverName: serverName,
|
||||||
|
serverID: serverID,
|
||||||
toolName: name,
|
toolName: name,
|
||||||
registeredName: registered,
|
registeredName: registered,
|
||||||
description: mcpTool.Description,
|
description: mcpTool.Description,
|
||||||
@@ -60,6 +66,7 @@ func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Poin
|
|||||||
clientPtr: clientPtr,
|
clientPtr: clientPtr,
|
||||||
timeoutSec: timeoutSec,
|
timeoutSec: timeoutSec,
|
||||||
connected: connected,
|
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) IsConnected() bool { return t.connected.Load() }
|
||||||
|
|
||||||
func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Result {
|
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() {
|
if !t.connected.Load() {
|
||||||
return tools.ErrorResult(fmt.Sprintf("MCP server %q is disconnected", t.serverName))
|
return tools.ErrorResult(fmt.Sprintf("MCP server %q is disconnected", t.serverName))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package mcp
|
|||||||
import (
|
import (
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -93,7 +94,7 @@ func TestBridgeToolNaming(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Without prefix → auto-derived from server name
|
// 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" {
|
if bt.Name() != "mcp_myserver__query" {
|
||||||
t.Errorf("expected name=mcp_myserver__query, got %s", bt.Name())
|
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
|
// 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" {
|
if bt2.Name() != "mcp_pg__query" {
|
||||||
t.Errorf("expected name=mcp_pg__query, got %s", bt2.Name())
|
t.Errorf("expected name=mcp_pg__query, got %s", bt2.Name())
|
||||||
}
|
}
|
||||||
@@ -114,13 +115,13 @@ func TestBridgeToolNaming(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// With mcp_ prefix → unchanged
|
// 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" {
|
if bt3.Name() != "mcp_pg__query" {
|
||||||
t.Errorf("expected name=mcp_pg__query, got %s", bt3.Name())
|
t.Errorf("expected name=mcp_pg__query, got %s", bt3.Name())
|
||||||
}
|
}
|
||||||
|
|
||||||
// Server name with hyphens → sanitized to underscores
|
// 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" {
|
if bt4.Name() != "mcp_my_server__query" {
|
||||||
t.Errorf("expected name=mcp_my_server__query, got %s", bt4.Name())
|
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
|
// DB-backed servers
|
||||||
store store.MCPServerStore
|
store store.MCPServerStore
|
||||||
|
|
||||||
|
// Grant checker for runtime grant verification (nil = skip check)
|
||||||
|
grantChecker GrantChecker
|
||||||
|
|
||||||
// Shared connection pool (nil = config-only mode)
|
// Shared connection pool (nil = config-only mode)
|
||||||
pool *Pool
|
pool *Pool
|
||||||
poolServers map[string]struct{} // server names acquired from pool (for cleanup)
|
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.
|
// NewManager creates a new MCP Manager.
|
||||||
func NewManager(registry *tools.Registry, opts ...ManagerOption) *Manager {
|
func NewManager(registry *tools.Registry, opts ...ManagerOption) *Manager {
|
||||||
m := &Manager{
|
m := &Manager{
|
||||||
@@ -167,7 +178,8 @@ func (m *Manager) Start(ctx context.Context) error {
|
|||||||
continue
|
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)
|
slog.Warn("mcp.server.connect_failed", "server", name, "error", err)
|
||||||
errs = append(errs, fmt.Sprintf("%s: %v", name, 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
|
// Pool mode: acquire shared connection, create per-agent BridgeTools
|
||||||
tid := store.TenantIDFromContext(ctx)
|
tid := store.TenantIDFromContext(ctx)
|
||||||
if err := m.connectViaPool(ctx, tid, srv.Name, srv.Transport, srv.Command,
|
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
|
return err
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
// Per-agent mode: create per-agent connection
|
// Per-agent mode: create per-agent connection
|
||||||
if err := m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
|
if err := m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
|
||||||
rs.args, rs.env, srv.URL, rs.headers,
|
rs.args, rs.env, srv.URL, rs.headers,
|
||||||
srv.ToolPrefix, srv.TimeoutSec); err != nil {
|
srv.ToolPrefix, srv.TimeoutSec, srv.ID); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -387,7 +399,7 @@ func (m *Manager) maybeEnterSearchMode() {
|
|||||||
|
|
||||||
// Update "mcp" group to only the kept inline names.
|
// Update "mcp" group to only the kept inline names.
|
||||||
inlineNames := allNames[:mcpToolInlineMaxCount]
|
inlineNames := allNames[:mcpToolInlineMaxCount]
|
||||||
tools.RegisterToolGroup("mcp", inlineNames)
|
m.registry.RegisterToolGroup("mcp", inlineNames)
|
||||||
m.searchMode = true
|
m.searchMode = true
|
||||||
|
|
||||||
slog.Info("mcp.search_mode.enabled",
|
slog.Info("mcp.search_mode.enabled",
|
||||||
@@ -459,7 +471,7 @@ func (m *Manager) ActivateTools(names []string) {
|
|||||||
}
|
}
|
||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
|
|
||||||
tools.RegisterToolGroup("mcp", activeNames)
|
m.registry.RegisterToolGroup("mcp", activeNames)
|
||||||
slog.Info("mcp.tools.activated", "tools", activated)
|
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).
|
// Register in registry outside lock (registry has its own sync).
|
||||||
m.registry.Register(bt)
|
m.registry.Register(bt)
|
||||||
tools.RegisterToolGroup("mcp", activeNames)
|
m.registry.RegisterToolGroup("mcp", activeNames)
|
||||||
slog.Info("mcp.tools.activated", "tools", []string{name})
|
slog.Info("mcp.tools.activated", "tools", []string{name})
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,7 +11,6 @@ import (
|
|||||||
mcpclient "github.com/mark3labs/mcp-go/client"
|
mcpclient "github.com/mark3labs/mcp-go/client"
|
||||||
"github.com/mark3labs/mcp-go/client/transport"
|
"github.com/mark3labs/mcp-go/client/transport"
|
||||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
// connectAndDiscover creates a client, initializes the MCP handshake, and
|
// 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.
|
// 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)
|
ss, mcpTools, err := connectAndDiscover(ctx, name, transportType, command, args, env, url, headers, timeoutSec)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Register tools
|
// Register tools
|
||||||
registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec)
|
registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec, serverID)
|
||||||
ss.toolNames = registeredNames
|
ss.toolNames = registeredNames
|
||||||
|
|
||||||
// Create health monitoring context
|
// Create health monitoring context
|
||||||
@@ -121,7 +121,7 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
|
|||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
|
|
||||||
if len(registeredNames) > 0 {
|
if len(registeredNames) > 0 {
|
||||||
tools.RegisterToolGroup("mcp:"+name, registeredNames)
|
m.registry.RegisterToolGroup("mcp:"+name, registeredNames)
|
||||||
m.updateMCPGroup()
|
m.updateMCPGroup()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -138,10 +138,11 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
|
|||||||
|
|
||||||
// registerBridgeTools creates BridgeTools from MCP tool definitions and
|
// registerBridgeTools creates BridgeTools from MCP tool definitions and
|
||||||
// registers them in the Manager's registry. Returns registered tool names.
|
// 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
|
var registeredNames []string
|
||||||
for _, mcpTool := range mcpTools {
|
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 {
|
if _, exists := m.registry.Get(bt.Name()); exists {
|
||||||
slog.Warn("mcp.tool.name_collision",
|
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
|
// connectViaPool acquires a shared connection from the pool and creates
|
||||||
// per-agent BridgeTools pointing to the shared client/connected pointers.
|
// 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)
|
entry, err := m.pool.Acquire(ctx, tenantID, name, transportType, command, args, env, url, headers, timeoutSec)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create per-agent BridgeTools from the pool's shared connection
|
// 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.
|
// Track server state and per-agent tool names.
|
||||||
// poolServers/poolToolNames keyed by plain name for Close() iteration.
|
// 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()
|
m.mu.Unlock()
|
||||||
|
|
||||||
if len(registeredNames) > 0 {
|
if len(registeredNames) > 0 {
|
||||||
tools.RegisterToolGroup("mcp:"+name, registeredNames)
|
m.registry.RegisterToolGroup("mcp:"+name, registeredNames)
|
||||||
m.updateMCPGroup()
|
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,
|
// registerPoolBridgeTools creates BridgeTools from pool entry's discovered tools,
|
||||||
// pointing to the shared client/connected pointers. Returns registered tool names.
|
// 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
|
var registeredNames []string
|
||||||
for _, mcpTool := range entry.tools {
|
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 {
|
if _, exists := m.registry.Get(bt.Name()); exists {
|
||||||
slog.Warn("mcp.tool.name_collision",
|
slog.Warn("mcp.tool.name_collision",
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||||
)
|
)
|
||||||
@@ -15,7 +16,7 @@ func makeBridgeTool(serverName, toolName string) *BridgeTool {
|
|||||||
Name: toolName,
|
Name: toolName,
|
||||||
Description: "test tool " + toolName,
|
Description: "test tool " + toolName,
|
||||||
InputSchema: mcpgo.ToolInputSchema{Type: "object"},
|
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
|
// setupSearchModeManager returns a Manager already in search mode with the given
|
||||||
|
|||||||
@@ -8,7 +8,6 @@ import (
|
|||||||
|
|
||||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
// UserCredServers returns servers requiring per-user credentials.
|
// UserCredServers returns servers requiring per-user credentials.
|
||||||
@@ -53,9 +52,9 @@ func (m *Manager) ServerToolNames(serverName string) []string {
|
|||||||
func (m *Manager) updateMCPGroup() {
|
func (m *Manager) updateMCPGroup() {
|
||||||
allNames := m.ToolNames()
|
allNames := m.ToolNames()
|
||||||
if len(allNames) > 0 {
|
if len(allNames) > 0 {
|
||||||
tools.RegisterToolGroup("mcp", allNames)
|
m.registry.RegisterToolGroup("mcp", allNames)
|
||||||
} else {
|
} else {
|
||||||
tools.UnregisterToolGroup("mcp")
|
m.registry.UnregisterToolGroup("mcp")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -88,7 +87,7 @@ func (m *Manager) unregisterAllTools() {
|
|||||||
m.registry.Unregister(toolName)
|
m.registry.Unregister(toolName)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
tools.UnregisterToolGroup("mcp:" + name)
|
m.registry.UnregisterToolGroup("mcp:" + name)
|
||||||
slog.Debug("mcp.server.unregistered", "server", name)
|
slog.Debug("mcp.server.unregistered", "server", name)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -105,7 +104,7 @@ func (m *Manager) unregisterAllTools() {
|
|||||||
m.servers = make(map[string]*serverState)
|
m.servers = make(map[string]*serverState)
|
||||||
m.poolServers = nil
|
m.poolServers = nil
|
||||||
m.poolToolNames = nil
|
m.poolToolNames = nil
|
||||||
tools.UnregisterToolGroup("mcp")
|
m.registry.UnregisterToolGroup("mcp")
|
||||||
}
|
}
|
||||||
|
|
||||||
// ToolInfo holds a tool's name and description for API responses.
|
// ToolInfo holds a tool's name and description for API responses.
|
||||||
|
|||||||
+96
-91
@@ -10,11 +10,9 @@ import (
|
|||||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||||
)
|
)
|
||||||
|
|
||||||
// toolGroupsMu protects toolGroups from concurrent access.
|
// builtinToolGroups is const-like seed data for per-Registry tool groups.
|
||||||
var toolGroupsMu sync.RWMutex
|
// Do NOT modify at runtime — each Registry gets a deep copy in NewRegistry().
|
||||||
|
var builtinToolGroups = map[string][]string{
|
||||||
// Tool groups map group names to tool names.
|
|
||||||
var toolGroups = map[string][]string{
|
|
||||||
"memory": {"memory_search", "memory_get"},
|
"memory": {"memory_search", "memory_get"},
|
||||||
"web": {"web_search", "web_fetch"},
|
"web": {"web_search", "web_fetch"},
|
||||||
"fs": {"read_file", "write_file", "list_files", "edit"},
|
"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.
|
// Package-level wrappers are REMOVED — use Registry methods instead.
|
||||||
// Used by the MCP manager to register "mcp" and "mcp:{serverName}" groups.
|
// See Registry.RegisterToolGroup, Registry.MergeToolGroup, Registry.UnregisterToolGroup.
|
||||||
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()
|
|
||||||
}
|
|
||||||
|
|
||||||
// Tool profiles define preset allow sets.
|
// Tool profiles define preset allow sets.
|
||||||
var toolProfiles = map[string][]string{
|
var toolProfiles = map[string][]string{
|
||||||
@@ -225,6 +192,11 @@ func (pe *PolicyEngine) evaluate(
|
|||||||
) []string {
|
) []string {
|
||||||
g := pe.globalPolicy
|
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
|
// Step 1: Global profile
|
||||||
allowed := pe.applyProfile(allTools, g.Profile)
|
allowed := pe.applyProfile(allTools, g.Profile)
|
||||||
|
|
||||||
@@ -237,49 +209,49 @@ func (pe *PolicyEngine) evaluate(
|
|||||||
|
|
||||||
// Step 3: Global allow list (restricts to only these)
|
// Step 3: Global allow list (restricts to only these)
|
||||||
if len(g.Allow) > 0 {
|
if len(g.Allow) > 0 {
|
||||||
allowed = intersectWithSpec(allowed, g.Allow)
|
allowed = intersectWithSpec(reg, allowed, g.Allow)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Step 4: Provider-level allow override
|
// Step 4: Provider-level allow override
|
||||||
if g.ByProvider != nil {
|
if g.ByProvider != nil {
|
||||||
if pp, ok := g.ByProvider[providerName]; ok && len(pp.Allow) > 0 {
|
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
|
// Step 5: Per-agent allow
|
||||||
if agentToolPolicy != nil && len(agentToolPolicy.Allow) > 0 {
|
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
|
// Step 6: Per-agent per-provider allow
|
||||||
if agentToolPolicy != nil && agentToolPolicy.ByProvider != nil {
|
if agentToolPolicy != nil && agentToolPolicy.ByProvider != nil {
|
||||||
if pp, ok := agentToolPolicy.ByProvider[providerName]; ok && len(pp.Allow) > 0 {
|
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
|
// Step 7: Group-level allow
|
||||||
if len(groupToolAllow) > 0 {
|
if len(groupToolAllow) > 0 {
|
||||||
allowed = intersectWithSpec(allowed, groupToolAllow)
|
allowed = intersectWithSpec(reg, allowed, groupToolAllow)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Apply global deny
|
// Apply global deny
|
||||||
if len(g.Deny) > 0 {
|
if len(g.Deny) > 0 {
|
||||||
allowed = subtractSpec(allowed, g.Deny)
|
allowed = subtractSpec(reg, allowed, g.Deny)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Apply agent deny
|
// Apply agent deny
|
||||||
if agentToolPolicy != nil && len(agentToolPolicy.Deny) > 0 {
|
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)
|
// Apply alsoAllow (additive — adds back tools without removing existing)
|
||||||
if len(g.AlsoAllow) > 0 {
|
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 {
|
if agentToolPolicy != nil && len(agentToolPolicy.AlsoAllow) > 0 {
|
||||||
allowed = unionWithSpec(allowed, allTools, agentToolPolicy.AlsoAllow)
|
allowed = unionWithSpec(reg, allowed, allTools, agentToolPolicy.AlsoAllow)
|
||||||
}
|
}
|
||||||
|
|
||||||
return allowed
|
return allowed
|
||||||
@@ -298,31 +270,34 @@ func (pe *PolicyEngine) applyProfile(allTools []string, profile string) []string
|
|||||||
return copySlice(allTools)
|
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,
|
// expandSpec expands a spec list (which may contain "group:xxx") into concrete tool names,
|
||||||
// filtered against available tools.
|
// filtered against available tools. Uses per-Registry tool groups to avoid cross-agent races.
|
||||||
func expandSpec(available []string, spec []string) []string {
|
func expandSpec(reg *Registry, available []string, spec []string) []string {
|
||||||
toolGroupsMu.RLock()
|
if reg == nil {
|
||||||
defer toolGroupsMu.RUnlock()
|
// 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)
|
expanded := make(map[string]bool)
|
||||||
for _, s := range spec {
|
for _, s := range spec {
|
||||||
if after, ok := strings.CutPrefix(s, "group:"); ok {
|
if !strings.HasPrefix(s, "group:") {
|
||||||
groupName := after
|
|
||||||
if members, ok := toolGroups[groupName]; ok {
|
|
||||||
for _, m := range members {
|
|
||||||
expanded[m] = true
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
expanded[s] = true
|
expanded[s] = true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
var result []string
|
var result []string
|
||||||
for _, t := range available {
|
for _, t := range available {
|
||||||
if expanded[t] {
|
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).
|
// intersectWithSpec keeps only tools in `current` that match the spec (with group expansion).
|
||||||
func intersectWithSpec(current []string, spec []string) []string {
|
func intersectWithSpec(reg *Registry, current []string, spec []string) []string {
|
||||||
toolGroupsMu.RLock()
|
if reg == nil {
|
||||||
defer toolGroupsMu.RUnlock()
|
return intersectWithSpecNoGroups(current, spec)
|
||||||
|
}
|
||||||
|
reg.toolGroupsMu.RLock()
|
||||||
|
defer reg.toolGroupsMu.RUnlock()
|
||||||
|
|
||||||
expanded := make(map[string]bool)
|
expanded := make(map[string]bool)
|
||||||
for _, s := range spec {
|
for _, s := range spec {
|
||||||
if after, ok := strings.CutPrefix(s, "group:"); ok {
|
if after, ok := strings.CutPrefix(s, "group:"); ok {
|
||||||
groupName := after
|
if members, ok := reg.toolGroups[after]; ok {
|
||||||
if members, ok := toolGroups[groupName]; ok {
|
|
||||||
for _, m := range members {
|
for _, m := range members {
|
||||||
expanded[m] = true
|
expanded[m] = true
|
||||||
}
|
}
|
||||||
@@ -360,16 +337,34 @@ func intersectWithSpec(current []string, spec []string) []string {
|
|||||||
return result
|
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.
|
// subtractSpec removes tools matching the spec (with group expansion) from current.
|
||||||
func subtractSpec(current []string, spec []string) []string {
|
func subtractSpec(reg *Registry, current []string, spec []string) []string {
|
||||||
toolGroupsMu.RLock()
|
if reg == nil {
|
||||||
defer toolGroupsMu.RUnlock()
|
return subtractSpecNoGroups(current, spec)
|
||||||
|
}
|
||||||
|
reg.toolGroupsMu.RLock()
|
||||||
|
defer reg.toolGroupsMu.RUnlock()
|
||||||
|
|
||||||
denied := make(map[string]bool)
|
denied := make(map[string]bool)
|
||||||
for _, s := range spec {
|
for _, s := range spec {
|
||||||
if after, ok := strings.CutPrefix(s, "group:"); ok {
|
if after, ok := strings.CutPrefix(s, "group:"); ok {
|
||||||
groupName := after
|
if members, ok := reg.toolGroups[after]; ok {
|
||||||
if members, ok := toolGroups[groupName]; ok {
|
|
||||||
for _, m := range members {
|
for _, m := range members {
|
||||||
denied[m] = true
|
denied[m] = true
|
||||||
}
|
}
|
||||||
@@ -388,6 +383,22 @@ func subtractSpec(current []string, spec []string) []string {
|
|||||||
return result
|
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.
|
// subtractSet removes exact tool names from current.
|
||||||
func subtractSet(current []string, deny []string) []string {
|
func subtractSet(current []string, deny []string) []string {
|
||||||
denied := make(map[string]bool, len(deny))
|
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.
|
// 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))
|
existing := make(map[string]bool, len(current))
|
||||||
for _, t := range current {
|
for _, t := range current {
|
||||||
existing[t] = true
|
existing[t] = true
|
||||||
}
|
}
|
||||||
|
|
||||||
toAdd := expandSpec(allTools, spec)
|
toAdd := expandSpec(reg, allTools, spec)
|
||||||
for _, t := range toAdd {
|
for _, t := range toAdd {
|
||||||
if !existing[t] {
|
if !existing[t] {
|
||||||
current = append(current, 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.
|
// 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.
|
// Used to prevent lazy-activated deferred tools from bypassing the deny list.
|
||||||
func (pe *PolicyEngine) IsDenied(name string, agentPolicy *config.ToolPolicySpec) bool {
|
func (pe *PolicyEngine) IsDenied(name string, agentPolicy *config.ToolPolicySpec) bool {
|
||||||
|
pe.mu.RLock()
|
||||||
|
reg := pe.registry
|
||||||
|
pe.mu.RUnlock()
|
||||||
|
|
||||||
if pe.globalPolicy != nil {
|
if pe.globalPolicy != nil {
|
||||||
if matchDenySpec(name, pe.globalPolicy.Deny) {
|
if matchDenySpec(reg, name, pe.globalPolicy.Deny) {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if agentPolicy != nil {
|
if agentPolicy != nil {
|
||||||
if matchDenySpec(name, agentPolicy.Deny) {
|
if matchDenySpec(reg, name, agentPolicy.Deny) {
|
||||||
return true
|
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).
|
// matchDenySpec returns true if name matches any entry in the deny spec (with group expansion).
|
||||||
func matchDenySpec(name string, spec []string) bool {
|
func matchDenySpec(reg *Registry, name string, spec []string) bool {
|
||||||
toolGroupsMu.RLock()
|
if reg == nil {
|
||||||
defer toolGroupsMu.RUnlock()
|
// No groups to expand — plain match only
|
||||||
|
return slices.Contains(spec, name)
|
||||||
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
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
return false
|
return reg.MatchDenySpec(name, spec)
|
||||||
}
|
}
|
||||||
|
|
||||||
// StripToolPrefix removes a prefix pattern from a tool name returned by the LLM.
|
// StripToolPrefix removes a prefix pattern from a tool name returned by the LLM.
|
||||||
|
|||||||
@@ -2,11 +2,13 @@ package tools
|
|||||||
|
|
||||||
import "testing"
|
import "testing"
|
||||||
|
|
||||||
func TestRegisterToolGroup(t *testing.T) {
|
func TestRegistry_RegisterToolGroup(t *testing.T) {
|
||||||
// Register a new MCP group
|
reg := NewRegistry()
|
||||||
RegisterToolGroup("mcp:postgres", []string{"mcp_pg__query", "mcp_pg__list_tables"})
|
|
||||||
|
|
||||||
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 {
|
if !ok {
|
||||||
t.Fatal("expected mcp:postgres group to exist")
|
t.Fatal("expected mcp:postgres group to exist")
|
||||||
}
|
}
|
||||||
@@ -15,25 +17,26 @@ func TestRegisterToolGroup(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Unregister
|
// Unregister
|
||||||
UnregisterToolGroup("mcp:postgres")
|
reg.UnregisterToolGroup("mcp:postgres")
|
||||||
if _, ok := toolGroups["mcp:postgres"]; ok {
|
if _, ok := reg.GetToolGroup("mcp:postgres"); ok {
|
||||||
t.Error("expected mcp:postgres group to be removed")
|
t.Error("expected mcp:postgres group to be removed")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestRegisterToolGroup_UsedInExpand(t *testing.T) {
|
func TestRegistry_RegisterToolGroup_UsedInExpand(t *testing.T) {
|
||||||
RegisterToolGroup("mcp:test", []string{"mcp_test__tool_a", "mcp_test__tool_b"})
|
reg := NewRegistry()
|
||||||
defer UnregisterToolGroup("mcp:test")
|
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"}
|
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 {
|
if len(expanded) != 2 {
|
||||||
t.Errorf("expected 2 tools from group:mcp:test, got %d: %v", len(expanded), expanded)
|
t.Errorf("expected 2 tools from group:mcp:test, got %d: %v", len(expanded), expanded)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Verify it works with subtractSpec too
|
// 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 {
|
if len(remaining) != 2 {
|
||||||
t.Errorf("expected 2 remaining after subtract, got %d: %v", len(remaining), remaining)
|
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
|
rateLimiter *ToolRateLimiter // nil = no rate limiting
|
||||||
scrubbing bool // scrub credentials from output (default true)
|
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
|
// 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.
|
// a deferred MCP tool. Returns true if the tool was successfully activated.
|
||||||
deferredActivator func(name string) bool
|
deferredActivator func(name string) bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewRegistry() *Registry {
|
func NewRegistry() *Registry {
|
||||||
return &Registry{
|
r := &Registry{
|
||||||
tools: make(map[string]Tool),
|
tools: make(map[string]Tool),
|
||||||
metadata: make(map[string]ToolMetadata),
|
metadata: make(map[string]ToolMetadata),
|
||||||
aliases: make(map[string]string),
|
aliases: make(map[string]string),
|
||||||
disabled: make(map[string]bool),
|
disabled: make(map[string]bool),
|
||||||
scrubbing: true, // enabled by default
|
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.
|
// SetDeferredActivator registers a callback that activates deferred tools on demand.
|
||||||
@@ -337,11 +348,15 @@ func (r *Registry) Count() int {
|
|||||||
func (r *Registry) Clone() *Registry {
|
func (r *Registry) Clone() *Registry {
|
||||||
r.mu.RLock()
|
r.mu.RLock()
|
||||||
defer r.mu.RUnlock()
|
defer r.mu.RUnlock()
|
||||||
|
r.toolGroupsMu.RLock()
|
||||||
|
defer r.toolGroupsMu.RUnlock()
|
||||||
|
|
||||||
clone := &Registry{
|
clone := &Registry{
|
||||||
tools: make(map[string]Tool, len(r.tools)),
|
tools: make(map[string]Tool, len(r.tools)),
|
||||||
metadata: make(map[string]ToolMetadata, len(r.metadata)),
|
metadata: make(map[string]ToolMetadata, len(r.metadata)),
|
||||||
aliases: make(map[string]string, len(r.aliases)),
|
aliases: make(map[string]string, len(r.aliases)),
|
||||||
disabled: make(map[string]bool, len(r.disabled)),
|
disabled: make(map[string]bool, len(r.disabled)),
|
||||||
|
toolGroups: make(map[string][]string, len(r.toolGroups)),
|
||||||
rateLimiter: r.rateLimiter,
|
rateLimiter: r.rateLimiter,
|
||||||
scrubbing: r.scrubbing,
|
scrubbing: r.scrubbing,
|
||||||
}
|
}
|
||||||
@@ -349,5 +364,102 @@ func (r *Registry) Clone() *Registry {
|
|||||||
maps.Copy(clone.metadata, r.metadata)
|
maps.Copy(clone.metadata, r.metadata)
|
||||||
maps.Copy(clone.aliases, r.aliases)
|
maps.Copy(clone.aliases, r.aliases)
|
||||||
maps.Copy(clone.disabled, r.disabled)
|
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
|
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