Skip to content

Commit 08198db

Browse files
claudeclaude
committed
fix: address code review issues C1, C2, I1
C1: Default /mcp endpoint now respects routing_mode config by using GetMCPServerForMode() instead of always using retrieve_tools server. C2: lookupToolPermission now uses exact match via lookupToolAnnotations (StateView) as primary lookup, with BM25 index search (limit 20) as fallback. Prevents wrong permission tier from fuzzy search results. I1: When code execution is disabled in config, the /mcp/code endpoint returns a clear error message instead of confusing runtime failure.
1 parent 2045c05 commit 08198db

3 files changed

Lines changed: 76 additions & 41 deletions

File tree

‎internal/server/mcp_code_execution.go‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -548,11 +548,21 @@ func (u *upstreamToolCaller) storeToolCallInHistory(serverName, toolName string,
548548
// lookupToolPermission returns the required permission tier for a tool based on its annotations.
549549
// This is used by the JS runtime to enforce auth context permissions during code_execution.
550550
func (p *MCPProxyServer) lookupToolPermission(serverName, toolName string) string {
551-
// Try to find tool annotations from the index
551+
// Primary: exact match via StateView (no BM25 fuzzy matching)
552+
annotations := p.lookupToolAnnotations(serverName, toolName)
553+
if annotations != nil {
554+
callWith := contracts.DeriveCallWith(annotations)
555+
perm := contracts.ToolVariantToOperationType[callWith]
556+
if perm != "" {
557+
return perm
558+
}
559+
}
560+
561+
// Fallback: search the index with enough candidates to find an exact match
552562
if p.index != nil {
553563
qualifiedName := serverName + ":" + toolName
554-
results, err := p.index.Search(qualifiedName, 1)
555-
if err == nil && len(results) > 0 {
564+
results, err := p.index.Search(qualifiedName, 20)
565+
if err == nil {
556566
for _, r := range results {
557567
if r.Tool != nil && r.Tool.ServerName == serverName && r.Tool.Name == toolName {
558568
callWith := contracts.DeriveCallWith(r.Tool.Annotations)
@@ -564,6 +574,7 @@ func (p *MCPProxyServer) lookupToolPermission(serverName, toolName string) strin
564574
}
565575
}
566576
}
577+
567578
// Default to read (safest)
568579
return contracts.OperationTypeRead
569580
}

‎internal/server/mcp_routing.go‎

Lines changed: 56 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -205,45 +205,64 @@ func (p *MCPProxyServer) makeDirectModeHandler(serverName, toolName string, anno
205205
// Includes: code_execution (with enhanced description listing available tools) and retrieve_tools (for discovery).
206206
// Does NOT include call_tool_read/write/destructive.
207207
func (p *MCPProxyServer) buildCodeExecModeTools() []mcpserver.ServerTool {
208-
ctx := context.Background()
209-
210-
// Build enhanced description with available tools catalog
211-
toolCatalog := p.buildToolCatalogDescription(ctx)
212-
213-
codeExecDescription := fmt.Sprintf(
214-
"Execute JavaScript code that orchestrates multiple upstream MCP tools in a single request. "+
215-
"Use this when you need to combine results from 2+ tools, implement conditional logic, loops, or data transformations.\n\n"+
216-
"**Available in JavaScript**:\n"+
217-
"- `input` global: Your input data passed via the 'input' parameter\n"+
218-
"- `call_tool(serverName, toolName, args)`: Call upstream tools (returns {ok, result} or {ok, error})\n"+
219-
"- Standard ES5.1+ JavaScript (no require(), filesystem, or network access)\n\n"+
220-
"**Available tools for orchestration**:\n%s\n\n"+
221-
"Use call_tool('serverName', 'toolName', {args}) to invoke tools.",
222-
toolCatalog,
223-
)
224-
225208
tools := make([]mcpserver.ServerTool, 0, 2)
226209

227-
// code_execution tool with enhanced description
228-
codeExecutionTool := mcp.NewTool("code_execution",
229-
mcp.WithDescription(codeExecDescription),
230-
mcp.WithTitleAnnotation("Code Execution"),
231-
mcp.WithDestructiveHintAnnotation(true),
232-
mcp.WithString("code",
233-
mcp.Required(),
234-
mcp.Description("JavaScript source code (ES5.1+) to execute. Use `input` to access input data and `call_tool(serverName, toolName, args)` to invoke upstream tools."),
235-
),
236-
mcp.WithObject("input",
237-
mcp.Description("Input data accessible as global `input` variable in JavaScript code (default: {})"),
238-
),
239-
mcp.WithObject("options",
240-
mcp.Description("Execution options: timeout_ms (1-600000, default: 120000), max_tool_calls (>= 0, 0=unlimited), allowed_servers (array of server names, empty=all allowed)"),
241-
),
242-
)
243-
tools = append(tools, mcpserver.ServerTool{
244-
Tool: codeExecutionTool,
245-
Handler: p.handleCodeExecution,
246-
})
210+
// Check if code execution is enabled in config
211+
if p.config != nil && !p.config.EnableCodeExecution {
212+
// Code execution is disabled: register a stub tool that returns a clear error message
213+
codeExecutionTool := mcp.NewTool("code_execution",
214+
mcp.WithDescription("Code execution is currently disabled. Enable it by setting \"enable_code_execution\": true in your mcpproxy config."),
215+
mcp.WithTitleAnnotation("Code Execution (Disabled)"),
216+
mcp.WithReadOnlyHintAnnotation(true),
217+
mcp.WithString("code",
218+
mcp.Required(),
219+
mcp.Description("JavaScript source code (ES5.1+) to execute."),
220+
),
221+
)
222+
tools = append(tools, mcpserver.ServerTool{
223+
Tool: codeExecutionTool,
224+
Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
225+
return mcp.NewToolResultError("Code execution is disabled. Enable it by setting \"enable_code_execution\": true in your mcpproxy configuration file."), nil
226+
},
227+
})
228+
} else {
229+
// Build enhanced description with available tools catalog
230+
ctx := context.Background()
231+
toolCatalog := p.buildToolCatalogDescription(ctx)
232+
233+
codeExecDescription := fmt.Sprintf(
234+
"Execute JavaScript code that orchestrates multiple upstream MCP tools in a single request. "+
235+
"Use this when you need to combine results from 2+ tools, implement conditional logic, loops, or data transformations.\n\n"+
236+
"**Available in JavaScript**:\n"+
237+
"- `input` global: Your input data passed via the 'input' parameter\n"+
238+
"- `call_tool(serverName, toolName, args)`: Call upstream tools (returns {ok, result} or {ok, error})\n"+
239+
"- Standard ES5.1+ JavaScript (no require(), filesystem, or network access)\n\n"+
240+
"**Available tools for orchestration**:\n%s\n\n"+
241+
"Use call_tool('serverName', 'toolName', {args}) to invoke tools.",
242+
toolCatalog,
243+
)
244+
245+
// code_execution tool with enhanced description
246+
codeExecutionTool := mcp.NewTool("code_execution",
247+
mcp.WithDescription(codeExecDescription),
248+
mcp.WithTitleAnnotation("Code Execution"),
249+
mcp.WithDestructiveHintAnnotation(true),
250+
mcp.WithString("code",
251+
mcp.Required(),
252+
mcp.Description("JavaScript source code (ES5.1+) to execute. Use `input` to access input data and `call_tool(serverName, toolName, args)` to invoke upstream tools."),
253+
),
254+
mcp.WithObject("input",
255+
mcp.Description("Input data accessible as global `input` variable in JavaScript code (default: {})"),
256+
),
257+
mcp.WithObject("options",
258+
mcp.Description("Execution options: timeout_ms (1-600000, default: 120000), max_tool_calls (>= 0, 0=unlimited), allowed_servers (array of server names, empty=all allowed)"),
259+
),
260+
)
261+
tools = append(tools, mcpserver.ServerTool{
262+
Tool: codeExecutionTool,
263+
Handler: p.handleCodeExecution,
264+
})
265+
}
247266

248267
// retrieve_tools for discovery
249268
retrieveToolsTool := mcp.NewTool("retrieve_tools",

‎internal/server/server.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -427,7 +427,12 @@ func (s *Server) Start(ctx context.Context) error {
427427
zap.String("listen", listenAddr))
428428

429429
// Create Streamable HTTP server with custom routing
430-
streamableServer := server.NewStreamableHTTPServer(s.mcpProxy.GetMCPServer())
430+
// Use the MCP server instance that corresponds to the configured routing_mode
431+
routingMode := ""
432+
if cfg != nil {
433+
routingMode = cfg.RoutingMode
434+
}
435+
streamableServer := server.NewStreamableHTTPServer(s.mcpProxy.GetMCPServerForMode(routingMode))
431436

432437
// Create custom HTTP server for handling multiple routes
433438
if err := s.startCustomHTTPServer(ctx, streamableServer); err != nil {

0 commit comments

Comments
 (0)