1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
// Copyright 2026 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
package aflow
import (
"context"
"errors"
"fmt"
"time"
"github.com/google/syzkaller/pkg/aflow/trajectory"
"github.com/google/syzkaller/pkg/osutil"
"github.com/modelcontextprotocol/go-sdk/mcp"
)
// Export all tools (and actions as tools) for tools/syz-mcp to serve them over MCP protocol.
var MCPTools = map[*mcp.Tool]MCPToolFunc{}
type MCPToolFunc func(ctx *Context, args map[string]any) (*mcp.CallToolResult, error)
func NewMCPContext(ctx context.Context, workdir string, cache *Cache, state map[string]any) *Context {
return &Context{
Context: ctx,
Workdir: osutil.Abs(workdir),
cache: cache,
state: state,
onEvent: func(span *trajectory.Span) error { return nil },
stubContext: stubContext{
timeNow: time.Now,
},
}
}
func registerMCPTool[State, Args, Results any](t *funcTool[State, Args, Results]) {
tool := &mcp.Tool{
Name: t.Name,
Description: t.Description,
InputSchema: mustSchemaFor[Args](),
OutputSchema: mustSchemaFor[Results](),
}
handler := func(ctx *Context, args map[string]any) (*mcp.CallToolResult, error) {
res, err := t.execute(ctx, args)
reply := &mcp.CallToolResult{
StructuredContent: res,
}
if err != nil {
if callErr := new(badCallError); !errors.As(err, &callErr) {
return nil, err
}
reply.SetError(err)
}
return reply, nil
}
registerMCP(tool, handler)
}
func registerMCPAction[Args, Results any](a *funcAction[Args, Results]) {
tool := &mcp.Tool{
Name: a.name,
Description: a.name,
InputSchema: mustSchemaFor[struct{}](),
OutputSchema: mustSchemaFor[struct{}](),
}
handler := func(ctx *Context, args map[string]any) (*mcp.CallToolResult, error) {
err := a.execute(ctx)
reply := &mcp.CallToolResult{
StructuredContent: struct{}{},
}
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
}
|
