Skip to content

Commit bf55136

Browse files
Copilotlpcox
Copilot
and
lpcox
authored
Refactor agent-filtered tool registration
Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>
1 parent f6aebb2 commit bf55136

4 files changed

Lines changed: 97 additions & 48 deletions

File tree

‎internal/server/agent_policy_enforce.go‎

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package server
22

33
import (
4+
"context"
45
"strings"
56

67
"github.com/github/gh-aw-mcpg/internal/logger"
@@ -64,30 +65,24 @@ func createAgentFilteredUnifiedServer(us *UnifiedServer, agentID string) *sdk.Se
6465
server := newSDKServer("awmg-unified", logTransport)
6566

6667
us.toolsMu.RLock()
67-
tools := make([]*ToolInfo, 0, len(us.tools))
68+
tools := make([]ToolInfo, 0, len(us.tools))
6869
for _, t := range us.tools {
69-
tools = append(tools, t)
70+
tools = append(tools, *t)
7071
}
7172
us.toolsMu.RUnlock()
7273

73-
registered := 0
74-
for _, toolInfo := range tools {
75-
backendID := toolInfo.BackendID
76-
unprefixed := strings.TrimPrefix(toolInfo.Name, backendID+"___")
77-
if !us.agentCanUseTool(agentID, backendID, unprefixed) {
78-
continue
79-
}
80-
if toolInfo.Handler == nil {
81-
continue
82-
}
83-
registerToolWithoutValidation(server, &sdk.Tool{
84-
Name: toolInfo.Name,
85-
Description: toolInfo.Description,
86-
InputSchema: toolInfo.InputSchema,
87-
Annotations: toolInfo.Annotations,
88-
}, toolInfo.Handler)
89-
registered++
90-
}
74+
registered := registerFilteredTools(
75+
server,
76+
tools,
77+
agentID,
78+
func(toolInfo ToolInfo) (string, string) {
79+
return toolInfo.BackendID, strings.TrimPrefix(toolInfo.Name, toolInfo.BackendID+"___")
80+
},
81+
us.agentCanUseTool,
82+
func(toolInfo ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) {
83+
return toolInfo.Handler
84+
},
85+
)
9186

9287
logger.LogInfo("client", "Built per-agent unified tool view: agent=%s, tools=%d",
9388
util.HashIdentifierForLog(agentID), registered)

‎internal/server/agent_policy_visibility_test.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,35 @@ func listToolsViaInMemory(t *testing.T, server *sdk.Server) []string {
3737
return names
3838
}
3939

40+
func TestRegisterFilteredTools(t *testing.T) {
41+
server := newSDKServer("filtered-tools-test", logTransport)
42+
handler := func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) {
43+
return &sdk.CallToolResult{}, nil, nil
44+
}
45+
tools := []ToolInfo{
46+
{Name: "allowed", BackendID: "github", InputSchema: map[string]interface{}{"type": "object"}},
47+
{Name: "denied", BackendID: "github", InputSchema: map[string]interface{}{"type": "object"}},
48+
{Name: "no-handler", BackendID: "github", InputSchema: map[string]interface{}{"type": "object"}},
49+
}
50+
51+
registered := registerFilteredTools(
52+
server,
53+
tools,
54+
"alice",
55+
func(tool ToolInfo) (string, string) { return tool.BackendID, tool.Name },
56+
func(_ string, _ string, toolName string) bool { return toolName != "denied" },
57+
func(tool ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) {
58+
if tool.Name == "no-handler" {
59+
return nil
60+
}
61+
return handler
62+
},
63+
)
64+
65+
assert.Equal(t, 1, registered)
66+
assert.ElementsMatch(t, []string{"allowed"}, listToolsViaInMemory(t, server))
67+
}
68+
4069
func agentVisibilityServer(t *testing.T) *UnifiedServer {
4170
t.Helper()
4271
cfg := &config.Config{

‎internal/server/routed.go‎

Lines changed: 21 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -49,35 +49,28 @@ func createAgentFilteredServer(unifiedServer *UnifiedServer, backendID, agentID
4949
logRouted.Printf("Creating filtered server for %s with %d tools", backendID, len(tools))
5050
logRouted.Printf("Backend %s has %d tools available", backendID, len(tools))
5151

52-
// Register each tool (without prefix) using the unified server's handlers
53-
for _, toolInfo := range tools {
54-
// Capture for closure
55-
toolNameCopy := toolInfo.Name
52+
// Register each tool (without prefix) using the unified server's handlers.
53+
registerFilteredTools(
54+
server,
55+
tools,
56+
agentID,
57+
func(toolInfo ToolInfo) (string, string) {
58+
return backendID, toolInfo.Name
59+
},
60+
unifiedServer.agentCanUseTool,
61+
func(toolInfo ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error) {
62+
handler := unifiedServer.GetToolHandler(backendID, toolInfo.Name)
63+
if handler == nil {
64+
logRouted.Printf("WARNING: No handler found for %s___%s", backendID, toolInfo.Name)
65+
return nil
66+
}
5667

57-
// Per-agent tool visibility: skip tools this agent's policy does not permit.
58-
if !unifiedServer.agentCanUseTool(agentID, backendID, toolNameCopy) {
59-
continue
60-
}
61-
62-
// Get the unified server's handler for this tool
63-
handler := unifiedServer.GetToolHandler(backendID, toolInfo.Name)
64-
if handler == nil {
65-
logRouted.Printf("WARNING: No handler found for %s___%s", backendID, toolInfo.Name)
66-
continue
67-
}
68-
69-
// Use registerToolWithoutValidation to bypass JSON Schema validation, allowing
70-
// InputSchema from backends using different JSON Schema versions (e.g., draft-07).
71-
registerToolWithoutValidation(server, &sdk.Tool{
72-
Name: toolInfo.Name, // Without prefix for the client
73-
Description: toolInfo.Description,
74-
InputSchema: toolInfo.InputSchema, // Include schema for clients
75-
Annotations: toolInfo.Annotations, // Preserve readOnly/destructive hints
76-
}, func(ctx context.Context, req *sdk.CallToolRequest, _ interface{}) (*sdk.CallToolResult, interface{}, error) {
77-
logRouted.Printf("[ROUTED] Calling unified handler for: %s", toolNameCopy)
78-
return handler(ctx, req, nil)
79-
})
80-
}
68+
return func(ctx context.Context, req *sdk.CallToolRequest, _ interface{}) (*sdk.CallToolResult, interface{}, error) {
69+
logRouted.Printf("[ROUTED] Calling unified handler for: %s", toolInfo.Name)
70+
return handler(ctx, req, nil)
71+
}
72+
},
73+
)
8174

8275
return server
8376
}

‎internal/server/tool_registry_helpers.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,38 @@ func registerToolWithoutValidation(server *sdk.Server, tool *sdk.Tool, handler f
5656
})
5757
}
5858

59+
// registerFilteredTools registers tools allowed for an agent and returns their count.
60+
func registerFilteredTools(
61+
server *sdk.Server,
62+
tools []ToolInfo,
63+
agentID string,
64+
toolIdentity func(ToolInfo) (backendID, toolName string),
65+
canUseTool func(agentID, backendID, toolName string) bool,
66+
handlerFor func(ToolInfo) func(context.Context, *sdk.CallToolRequest, interface{}) (*sdk.CallToolResult, interface{}, error),
67+
) int {
68+
registered := 0
69+
for _, toolInfo := range tools {
70+
backendID, toolName := toolIdentity(toolInfo)
71+
if !canUseTool(agentID, backendID, toolName) {
72+
continue
73+
}
74+
75+
handler := handlerFor(toolInfo)
76+
if handler == nil {
77+
continue
78+
}
79+
80+
registerToolWithoutValidation(server, &sdk.Tool{
81+
Name: toolInfo.Name,
82+
Description: toolInfo.Description,
83+
InputSchema: toolInfo.InputSchema,
84+
Annotations: toolInfo.Annotations,
85+
}, handler)
86+
registered++
87+
}
88+
return registered
89+
}
90+
5991
func getToolResponseFilter(cfg *config.Config, serverID, toolName string) string {
6092
if cfg == nil {
6193
return ""

0 commit comments

Comments
 (0)