mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 06:17:41 +00:00
fix(agent): surface per-user MCP tools to the LLM tool list [B24:2794] (#1235)
require_user_credentials MCP servers loaded and connected, but their tool defs were never sent to the model, so the agent could never call the per-user mcp_<prefix>__* tools (it only ever saw the shared MCP server). The per-user tool objects are intentionally kept out of the shared registry to prevent cross-user credential leaks, so FilterTools could not emit them. Surface them through a request-scoped overlay registry passed to the policy engine: per-user tools are now evaluated AND emitted under the same allow/deny rules as registry tools, then resolved per-actor at execution time via executeToolForActor. The overlay is discarded after filtering, so the shared registry and other users stay unaffected. - add tools.NewUserToolOverlay (request-scoped, discarded post-filter) - buildFilteredTools: wrap registry in overlay when per-user tools present; no-policy path appends their defs directly - fix stale getUserMCPTools docstring (claimed shared-registry registration) - tests: overlay policy allow/deny + per-user surfacing / strip / dedup Co-authored-by: DangTinh311 <dangtinh31193@gmail.com>
This commit is contained in:
1 parent
719d5934bb
commit
8eef0509ce
7 files changed
+323
-13
No files matched your search
@@ -55,7 +55,7 @@ func TestImageGenGate_AllTrue_ToolPresent(t *testing.T) {
|
||||
prov := &imageCapableProvider{imageGen: true}
|
||||
l := buildImageGenLoop(true, prov)
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil)
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, nil)
|
||||
|
||||
if !hasImageGenTool(defs) {
|
||||
t.Error("expected image_generation tool when all gate conditions are true")
|
||||
@@ -68,7 +68,7 @@ func TestImageGenGate_ProviderNoCapability_ToolAbsent(t *testing.T) {
|
||||
prov := &imageCapableProvider{imageGen: false}
|
||||
l := buildImageGenLoop(true, prov)
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil)
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, nil)
|
||||
|
||||
if hasImageGenTool(defs) {
|
||||
t.Error("image_generation must NOT be in tools when provider does not advertise ImageGeneration")
|
||||
@@ -82,7 +82,7 @@ func TestImageGenGate_ProviderNotCapabilitiesAware_ToolAbsent(t *testing.T) {
|
||||
prov := &stubProvider{}
|
||||
l := buildImageGenLoop(true, prov)
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil)
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, nil)
|
||||
|
||||
if hasImageGenTool(defs) {
|
||||
t.Error("image_generation must NOT be in tools when provider is not CapabilitiesAware")
|
||||
@@ -95,7 +95,7 @@ func TestImageGenGate_AgentConfigDisabled_ToolAbsent(t *testing.T) {
|
||||
prov := &imageCapableProvider{imageGen: true}
|
||||
l := buildImageGenLoop(false, prov) // allowImageGeneration = false
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil)
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, nil)
|
||||
|
||||
if hasImageGenTool(defs) {
|
||||
t.Error("image_generation must NOT be in tools when agent config disables it")
|
||||
@@ -109,7 +109,7 @@ func TestImageGenGate_FinalIteration_AllToolsStripped(t *testing.T) {
|
||||
l := buildImageGenLoop(true, prov)
|
||||
|
||||
// iteration == maxIter → final stripping path; gate never reached
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 5, 5, nil)
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 5, 5, nil, nil)
|
||||
|
||||
if len(defs) != 0 {
|
||||
t.Errorf("final iteration must strip all tools; got %d: %v", len(defs), defs)
|
||||
|
||||
@@ -85,9 +85,12 @@ func resolveActorUserID(userID, senderID, peerKind, channelType string) string {
|
||||
}
|
||||
|
||||
// getUserMCPTools returns per-user MCP tools for servers requiring user credentials.
|
||||
// Tools are cached per-user in mcpUserTools sync.Map and registered in the shared
|
||||
// tool registry so ExecuteWithContext can resolve them. On first call for a user,
|
||||
// connections are established via pool.AcquireUser() and BridgeTools created.
|
||||
// Tools are cached per-user in mcpUserTools sync.Map and their NAMES added to the "mcp"
|
||||
// tool group; the tool objects are deliberately NOT registered in the shared registry
|
||||
// (cross-user identity leak — see inline note below). The returned slice is passed to
|
||||
// buildFilteredTools so the defs surface to the LLM; execution resolves per-actor via
|
||||
// executeToolForActor → mcpUserTools. On first call for a user, connections are
|
||||
// established via pool.AcquireUser() and BridgeTools created.
|
||||
func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool {
|
||||
if len(l.mcpUserCredSrvs) == 0 || l.mcpPool == nil || l.mcpStore == nil || userID == "" {
|
||||
if userID == "" && len(l.mcpUserCredSrvs) > 0 {
|
||||
|
||||
@@ -228,7 +228,7 @@ func (l *Loop) makeBuildFilteredTools(req *RunRequest) func(state *pipeline.RunS
|
||||
}
|
||||
allMsgs := state.Messages.All()
|
||||
toolDefs, _, returnedMsgs := l.buildFilteredTools(req, state.Context.HadBootstrap,
|
||||
state.Iteration, maxIter, allMsgs)
|
||||
state.Iteration, maxIter, allMsgs, userTools)
|
||||
// buildFilteredTools returns the full messages slice; only messages appended
|
||||
// beyond the original length are injections (e.g. final-iteration hint).
|
||||
// Appending the entire slice would duplicate system+history into pending.
|
||||
|
||||
@@ -15,22 +15,51 @@ var imageGenToolDef = providers.ToolDefinition{Type: "image_generation"}
|
||||
|
||||
// buildFilteredTools resolves the per-iteration tool definitions based on policy,
|
||||
// disabled tools, bootstrap mode, skill visibility, channel type, and iteration budget.
|
||||
// Per-user MCP tools must be registered in the Registry before calling this function
|
||||
// (via getUserMCPTools) so they are included in policy filtering and execution.
|
||||
// Per-user MCP tools (require_user_credentials servers) are passed in via userTools —
|
||||
// their objects deliberately live ONLY in l.mcpUserTools (cross-user isolation), NOT
|
||||
// in the shared registry. They are surfaced via a request-scoped overlay registry so
|
||||
// the PolicyEngine evaluates AND emits them under the SAME allow/deny rules as
|
||||
// registry tools; execution routes per-actor via executeToolForActor.
|
||||
// Returns tool definitions for the provider, an allowed-tools map for execution validation,
|
||||
// and the (potentially modified) messages slice when final-iteration stripping appends a hint.
|
||||
func (l *Loop) buildFilteredTools(req *RunRequest, hadBootstrap bool, iteration, maxIter int, messages []providers.Message) ([]providers.ToolDefinition, map[string]bool, []providers.Message) {
|
||||
func (l *Loop) buildFilteredTools(req *RunRequest, hadBootstrap bool, iteration, maxIter int, messages []providers.Message, userTools []tools.Tool) ([]providers.ToolDefinition, map[string]bool, []providers.Message) {
|
||||
// Build provider request with policy-filtered tools.
|
||||
var toolDefs []providers.ToolDefinition
|
||||
var allowedTools map[string]bool
|
||||
if l.toolPolicy != nil {
|
||||
toolDefs = l.toolPolicy.FilterTools(l.tools, l.id, l.provider.Name(), l.agentToolPolicy, req.ToolAllow, false, false)
|
||||
// Per-user MCP tool objects are NOT in the shared registry (cross-user
|
||||
// isolation — see getUserMCPTools), so FilterTools alone cannot emit them.
|
||||
// Wrap the registry in a request-scoped overlay that adds the calling actor's
|
||||
// per-user tools: FilterTools then evaluates them through the same policy
|
||||
// pipeline (profile/allow/deny, incl. an explicit deny on a per-user tool name)
|
||||
// and emits them via overlay.Get. The overlay is local and discarded after this
|
||||
// call, so the shared registry — and thus other users — stay unaffected.
|
||||
// NewUserToolOverlay returns l.tools unchanged when userTools is empty.
|
||||
registry := tools.NewUserToolOverlay(l.tools, userTools)
|
||||
toolDefs = l.toolPolicy.FilterTools(registry, l.id, l.provider.Name(), l.agentToolPolicy, req.ToolAllow, false, false)
|
||||
allowedTools = make(map[string]bool, len(toolDefs))
|
||||
for _, td := range toolDefs {
|
||||
allowedTools[td.Function.Name] = true
|
||||
}
|
||||
} else {
|
||||
// No policy → all tools allowed. ProviderDefs() omits per-user MCP tools (not in
|
||||
// the shared registry), so append their defs directly (dedup by name).
|
||||
toolDefs = l.tools.ProviderDefs()
|
||||
if len(userTools) > 0 {
|
||||
seen := make(map[string]bool, len(toolDefs))
|
||||
for _, td := range toolDefs {
|
||||
if td.Function != nil {
|
||||
seen[td.Function.Name] = true
|
||||
}
|
||||
}
|
||||
for _, t := range userTools {
|
||||
if t == nil || seen[t.Name()] {
|
||||
continue
|
||||
}
|
||||
seen[t.Name()] = true
|
||||
toolDefs = append(toolDefs, tools.ToProviderDef(t))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// V3 orchestration mode filtering: hide tools the agent shouldn't see.
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
package agent
|
||||
|
||||
// Regression tests: per-user MCP tools (require_user_credentials servers) must be
|
||||
// surfaced to the LLM tool list. Their objects live ONLY in l.mcpUserTools (cross-user
|
||||
// isolation) and are NEVER in the shared registry, so the ONLY path that exposes their
|
||||
// defs to the model is the userTools argument of buildFilteredTools. A prior regression
|
||||
// (anti-leak hardening) dropped that merge → the model could never call mcp_<prefix>__*.
|
||||
// See plans/reports/bugreport-260617-1708-per-user-mcp-tools-not-surfaced-to-llm.md.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
)
|
||||
|
||||
// fakeUserMCPTool is a minimal tools.Tool standing in for a per-user MCP BridgeTool.
|
||||
type fakeUserMCPTool struct{ name string }
|
||||
|
||||
func (f *fakeUserMCPTool) Name() string { return f.name }
|
||||
func (f *fakeUserMCPTool) Description() string { return "per-user mcp tool stub" }
|
||||
func (f *fakeUserMCPTool) Parameters() map[string]any { return map[string]any{"type": "object"} }
|
||||
func (f *fakeUserMCPTool) Execute(context.Context, map[string]any) *tools.Result {
|
||||
return &tools.Result{ForLLM: "ok"}
|
||||
}
|
||||
|
||||
// hasToolNamed reports whether defs contains a function tool with the given name.
|
||||
func hasToolNamed(defs []providers.ToolDefinition, name string) bool {
|
||||
for _, d := range defs {
|
||||
if d.Function != nil && d.Function.Name == name {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ─── Per-user MCP tools appear in the LLM tool list ───────────────────────
|
||||
|
||||
func TestUserMCPTools_SurfacedToLLM(t *testing.T) {
|
||||
l := buildImageGenLoop(false, &stubProvider{}) // toolPolicy nil → defs come only from userTools
|
||||
userTools := []tools.Tool{
|
||||
&fakeUserMCPTool{name: "mcp_bx24__search"},
|
||||
&fakeUserMCPTool{name: "mcp_bx24__execute"},
|
||||
}
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, userTools)
|
||||
|
||||
if !hasToolNamed(defs, "mcp_bx24__search") || !hasToolNamed(defs, "mcp_bx24__execute") {
|
||||
t.Errorf("per-user MCP tools must surface to the LLM tool list; got %d defs: %v", len(defs), defs)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Final iteration strips per-user MCP tools too (force text response) ───
|
||||
|
||||
func TestUserMCPTools_FinalIterationStripped(t *testing.T) {
|
||||
l := buildImageGenLoop(false, &stubProvider{})
|
||||
userTools := []tools.Tool{&fakeUserMCPTool{name: "mcp_bx24__execute"}}
|
||||
|
||||
// iteration == maxIter → final stripping path: all tools removed, including per-user.
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 5, 5, nil, userTools)
|
||||
|
||||
if len(defs) != 0 {
|
||||
t.Errorf("final iteration must strip per-user MCP tools; got %d: %v", len(defs), defs)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Duplicate per-user tool names emit only one def ──────────────────────
|
||||
|
||||
func TestUserMCPTools_DedupByName(t *testing.T) {
|
||||
l := buildImageGenLoop(false, &stubProvider{})
|
||||
userTools := []tools.Tool{
|
||||
&fakeUserMCPTool{name: "mcp_bx24__execute"},
|
||||
&fakeUserMCPTool{name: "mcp_bx24__execute"}, // duplicate
|
||||
}
|
||||
|
||||
defs, _, _ := l.buildFilteredTools(&RunRequest{}, false, 1, 10, nil, userTools)
|
||||
|
||||
count := 0
|
||||
for _, d := range defs {
|
||||
if d.Function != nil && d.Function.Name == "mcp_bx24__execute" {
|
||||
count++
|
||||
}
|
||||
}
|
||||
if count != 1 {
|
||||
t.Errorf("duplicate per-user tool name must emit exactly one def; got %d", count)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package tools
|
||||
|
||||
// userToolOverlay is a request-scoped ToolExecutor that overlays a set of per-user
|
||||
// tools on top of a base executor (the shared registry). It lets the PolicyEngine
|
||||
// evaluate AND emit per-user MCP tools (require_user_credentials servers) under the
|
||||
// SAME allow/deny rules as registry tools, WITHOUT registering the per-user tool
|
||||
// objects into the shared registry — doing so would leak one user's credentialed
|
||||
// BridgeTool to every other user (see agent.getUserMCPTools).
|
||||
//
|
||||
// Construct one per request from the calling actor's tools (see
|
||||
// agent.buildFilteredTools) and discard it after the tool-filtering pass. The base
|
||||
// executor is embedded, so every ToolExecutor method except List/Get delegates to
|
||||
// it unchanged.
|
||||
type userToolOverlay struct {
|
||||
ToolExecutor // embedded base (shared registry)
|
||||
extra map[string]Tool // per-user tools by name; take precedence in Get
|
||||
names []string // per-user names, insertion order, appended by List
|
||||
}
|
||||
|
||||
// NewUserToolOverlay returns a ToolExecutor exposing base's tools plus userTools.
|
||||
// Duplicate names within userTools collapse (first wins); nil entries are skipped.
|
||||
// For List(), per-user names already present in base are not duplicated. For Get(),
|
||||
// a per-user tool shadows a same-named base tool — per-user objects carry the
|
||||
// actor's credentials. Returns base unchanged when userTools is empty, so callers
|
||||
// pay nothing on the common no-per-user path.
|
||||
func NewUserToolOverlay(base ToolExecutor, userTools []Tool) ToolExecutor {
|
||||
if len(userTools) == 0 {
|
||||
return base
|
||||
}
|
||||
extra := make(map[string]Tool, len(userTools))
|
||||
names := make([]string, 0, len(userTools))
|
||||
for _, t := range userTools {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
name := t.Name()
|
||||
if _, dup := extra[name]; dup {
|
||||
continue
|
||||
}
|
||||
extra[name] = t
|
||||
names = append(names, name)
|
||||
}
|
||||
if len(extra) == 0 {
|
||||
return base
|
||||
}
|
||||
return &userToolOverlay{ToolExecutor: base, extra: extra, names: names}
|
||||
}
|
||||
|
||||
// Get resolves per-user tools first, then falls back to the base executor.
|
||||
func (o *userToolOverlay) Get(name string) (Tool, bool) {
|
||||
if t, ok := o.extra[name]; ok {
|
||||
return t, true
|
||||
}
|
||||
return o.ToolExecutor.Get(name)
|
||||
}
|
||||
|
||||
// List returns base tool names plus any per-user names not already present in base.
|
||||
func (o *userToolOverlay) List() []string {
|
||||
base := o.ToolExecutor.List()
|
||||
seen := make(map[string]bool, len(base))
|
||||
out := make([]string, 0, len(base)+len(o.names))
|
||||
for _, n := range base {
|
||||
seen[n] = true
|
||||
out = append(out, n)
|
||||
}
|
||||
for _, n := range o.names {
|
||||
if !seen[n] {
|
||||
out = append(out, n)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Compile-time check: the overlay satisfies ToolExecutor.
|
||||
var _ ToolExecutor = (*userToolOverlay)(nil)
|
||||
@@ -0,0 +1,115 @@
|
||||
package tools
|
||||
|
||||
// Tests for userToolOverlay — the request-scoped registry overlay that surfaces
|
||||
// per-user MCP tools through the PolicyEngine without registering them in the shared
|
||||
// registry. The decisive case is TestUserToolOverlay_PolicyDenyByNameSuppresses:
|
||||
// it proves per-user MCP tools now honor an agent/global deny (Finding 1 fix), which
|
||||
// the previous "append after FilterTools" approach bypassed.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
)
|
||||
|
||||
// overlayFakeTool is a minimal Tool standing in for a per-user MCP BridgeTool.
|
||||
type overlayFakeTool struct{ name string }
|
||||
|
||||
func (f *overlayFakeTool) Name() string { return f.name }
|
||||
func (f *overlayFakeTool) Description() string { return "fake " + f.name }
|
||||
func (f *overlayFakeTool) Parameters() map[string]any { return map[string]any{"type": "object"} }
|
||||
func (f *overlayFakeTool) Execute(context.Context, map[string]any) *Result {
|
||||
return &Result{ForLLM: "ok"}
|
||||
}
|
||||
|
||||
// emptyBaseExecutor is a ToolExecutor with no tools — stands in for a shared registry
|
||||
// that does NOT contain the per-user MCP tools (the real-world condition).
|
||||
type emptyBaseExecutor struct{}
|
||||
|
||||
func (emptyBaseExecutor) ExecuteWithContext(context.Context, string, map[string]any, string, string, string, string, AsyncCallback) *Result {
|
||||
return &Result{}
|
||||
}
|
||||
func (emptyBaseExecutor) TryActivateDeferred(string) bool { return false }
|
||||
func (emptyBaseExecutor) ProviderDefs() []providers.ToolDefinition { return nil }
|
||||
func (emptyBaseExecutor) Get(string) (Tool, bool) { return nil, false }
|
||||
func (emptyBaseExecutor) List() []string { return nil }
|
||||
func (emptyBaseExecutor) Aliases() map[string]string { return nil }
|
||||
|
||||
func defNameSet(defs []providers.ToolDefinition) map[string]bool {
|
||||
m := make(map[string]bool, len(defs))
|
||||
for _, d := range defs {
|
||||
if d.Function != nil {
|
||||
m[d.Function.Name] = true
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// ─── Overlay exposes per-user tools via List/Get ──────────────────────────
|
||||
|
||||
func TestUserToolOverlay_AddsPerUserTools(t *testing.T) {
|
||||
ov := NewUserToolOverlay(emptyBaseExecutor{}, []Tool{
|
||||
&overlayFakeTool{name: "mcp_bx24__execute"},
|
||||
&overlayFakeTool{name: "mcp_bx24__search"},
|
||||
})
|
||||
|
||||
list := ov.List()
|
||||
if len(list) != 2 {
|
||||
t.Fatalf("expected 2 names in overlay List, got %d: %v", len(list), list)
|
||||
}
|
||||
if _, ok := ov.Get("mcp_bx24__execute"); !ok {
|
||||
t.Error("overlay must resolve per-user tool via Get")
|
||||
}
|
||||
if _, ok := ov.Get("does_not_exist"); ok {
|
||||
t.Error("overlay must not resolve unknown tools")
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Empty / all-nil userTools returns base unchanged ─────────────────────
|
||||
|
||||
func TestUserToolOverlay_EmptyReturnsBase(t *testing.T) {
|
||||
base := emptyBaseExecutor{}
|
||||
if ov := NewUserToolOverlay(base, nil); ov.List() != nil {
|
||||
t.Error("nil userTools should return base unchanged (nil List)")
|
||||
}
|
||||
if ov := NewUserToolOverlay(base, []Tool{nil, nil}); ov.List() != nil {
|
||||
t.Error("all-nil userTools should return base unchanged (nil List)")
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Policy with default (full) profile emits per-user tools ──────────────
|
||||
|
||||
func TestUserToolOverlay_PolicyEmitsByDefault(t *testing.T) {
|
||||
pe := NewPolicyEngine(&config.ToolsConfig{}) // empty profile = full = all allowed
|
||||
ov := NewUserToolOverlay(emptyBaseExecutor{}, []Tool{
|
||||
&overlayFakeTool{name: "mcp_bx24__execute"},
|
||||
&overlayFakeTool{name: "mcp_bx24__search"},
|
||||
})
|
||||
|
||||
defs := pe.FilterTools(ov, "agent", "openai", nil, nil, false, false)
|
||||
names := defNameSet(defs)
|
||||
if !names["mcp_bx24__execute"] || !names["mcp_bx24__search"] {
|
||||
t.Errorf("full profile must emit per-user tools through the overlay; got %v", names)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Finding 1 fix: an explicit deny on a per-user tool name suppresses it ─
|
||||
|
||||
func TestUserToolOverlay_PolicyDenyByNameSuppresses(t *testing.T) {
|
||||
pe := NewPolicyEngine(&config.ToolsConfig{Deny: []string{"mcp_bx24__execute"}})
|
||||
ov := NewUserToolOverlay(emptyBaseExecutor{}, []Tool{
|
||||
&overlayFakeTool{name: "mcp_bx24__execute"},
|
||||
&overlayFakeTool{name: "mcp_bx24__search"},
|
||||
})
|
||||
|
||||
defs := pe.FilterTools(ov, "agent", "openai", nil, nil, false, false)
|
||||
names := defNameSet(defs)
|
||||
if names["mcp_bx24__execute"] {
|
||||
t.Error("denied per-user tool must NOT be emitted (Finding 1: policy must apply)")
|
||||
}
|
||||
if !names["mcp_bx24__search"] {
|
||||
t.Error("non-denied per-user tool must still be emitted")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user