From b19feb243f1f2bdb7f04e5c6f3b693ee8cd2a8b7 Mon Sep 17 00:00:00 2001
From: Dmitry Vyukov <dvyukov@google.com>
Date: Tue, 10 Mar 2026 16:09:13 +0100
Subject: pkg/aflow: ensure we don't register MCP tools with duplicate names

If we have duplicate names, then only one of the duplicates will be used at random.
Add a check that we don't have duplicate names.
Currently it's only "crash-reproducer" (both action and a tool).
Also ignore "set-results" tool, and all tools created in tests.
---
 pkg/aflow/flow/flows_test.go        | 14 +++++++++++++-
 pkg/aflow/flow_test.go              |  5 +++++
 pkg/aflow/llm_agent.go              |  4 +++-
 pkg/aflow/mcp.go                    | 19 ++++++++++++++++++-
 pkg/aflow/tool/syzlang/reproduce.go |  2 +-
 5 files changed, 40 insertions(+), 4 deletions(-)

(limited to 'pkg')

diff --git a/pkg/aflow/flow/flows_test.go b/pkg/aflow/flow/flows_test.go
index 3c01596b5..a66479d14 100644
--- a/pkg/aflow/flow/flows_test.go
+++ b/pkg/aflow/flow/flows_test.go
@@ -3,4 +3,16 @@
 
 package flow
 
-// An empty test that runs registration and verification of all registered workflows via init functions.
+import (
+	"testing"
+
+	"github.com/google/syzkaller/pkg/aflow"
+)
+
+// Note: this test also runs registration and verification of all registered workflows via init functions.
+
+func TestMCPTools(t *testing.T) {
+	for tool := range aflow.MCPTools {
+		t.Log(tool.Name)
+	}
+}
diff --git a/pkg/aflow/flow_test.go b/pkg/aflow/flow_test.go
index 734781c41..94758c549 100644
--- a/pkg/aflow/flow_test.go
+++ b/pkg/aflow/flow_test.go
@@ -15,6 +15,11 @@ import (
 	"google.golang.org/genai"
 )
 
+func init() {
+	// Tests register tools with duplicate names.
+	registerMCPTools = false
+}
+
 func TestWorkflow(t *testing.T) {
 	type flowInputs struct {
 		InFoo int
diff --git a/pkg/aflow/llm_agent.go b/pkg/aflow/llm_agent.go
index 378f044a5..3529fa5b1 100644
--- a/pkg/aflow/llm_agent.go
+++ b/pkg/aflow/llm_agent.go
@@ -108,7 +108,7 @@ func Tools(tools ...any) []Tool {
 // LLMOutputs creates a special tool that can be used by LLM to provide structured outputs.
 func LLMOutputs[Args any]() *llmOutputs {
 	return &llmOutputs{
-		tool: NewFuncTool("set-results", func(ctx *Context, state struct{}, args Args) (Args, error) {
+		tool: NewFuncTool(llmSetResultsTool, func(ctx *Context, state struct{}, args Args) (Args, error) {
 			return args, nil
 		}, "Use this tool to provide results of the analysis."),
 		provideOutputs: func(ctx *verifyContext, who string, many bool) {
@@ -129,6 +129,8 @@ func LLMOutputs[Args any]() *llmOutputs {
 	}
 }
 
+const llmSetResultsTool = "set-results"
+
 const llmOutputsInstruction = `
 
 Use set-results tool to provide results of the analysis.
diff --git a/pkg/aflow/mcp.go b/pkg/aflow/mcp.go
index f94376113..5e0944e08 100644
--- a/pkg/aflow/mcp.go
+++ b/pkg/aflow/mcp.go
@@ -6,6 +6,7 @@ package aflow
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/google/syzkaller/pkg/aflow/trajectory"
@@ -51,7 +52,7 @@ func registerMCPTool[State, Args, Results any](t *funcTool[State, Args, Results]
 		}
 		return reply, nil
 	}
-	MCPTools[tool] = handler
+	registerMCP(tool, handler)
 }
 
 func registerMCPAction[Args, Results any](a *funcAction[Args, Results]) {
@@ -68,5 +69,21 @@ func registerMCPAction[Args, Results any](a *funcAction[Args, Results]) {
 		}
 		return reply, err
 	}
+	registerMCP(tool, handler)
+}
+
+var (
+	registerMCPTools = true
+	mcpToolNames     = map[string]bool{}
+)
+
+func registerMCP(tool *mcp.Tool, handler MCPToolFunc) {
+	if !registerMCPTools || tool.Name == llmSetResultsTool {
+		return
+	}
+	if mcpToolNames[tool.Name] {
+		panic(fmt.Sprintf("MCP tool %q is already registered", tool.Name))
+	}
+	mcpToolNames[tool.Name] = true
 	MCPTools[tool] = handler
 }
diff --git a/pkg/aflow/tool/syzlang/reproduce.go b/pkg/aflow/tool/syzlang/reproduce.go
index 3faa11563..2c22bddd9 100644
--- a/pkg/aflow/tool/syzlang/reproduce.go
+++ b/pkg/aflow/tool/syzlang/reproduce.go
@@ -9,7 +9,7 @@ import (
 	_ "github.com/google/syzkaller/sys"
 )
 
-var Reproduce = aflow.NewFuncTool("crash-reproducer", reproduce, `
+var Reproduce = aflow.NewFuncTool("reproduce-crash", reproduce, `
 Tool evaluates whether the given syz repro program crashes the kernel.
 It will compile the program and execute it in a VM.
 `)
-- 
cgit mrf-deployment