Skip to content

Commit 5a92f9c

Browse files
omgitsadsomgitsads
committed
Add intermediate structs for toolsets
1 parent 8084335 commit 5a92f9c

2 files changed

Lines changed: 51 additions & 37 deletions

File tree

‎pkg/github/tools.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
"github.com/github/github-mcp-server/pkg/toolsets"
1010
"github.com/github/github-mcp-server/pkg/translations"
1111
"github.com/google/go-github/v77/github"
12-
"github.com/mark3labs/mcp-go/server"
12+
"github.com/modelcontextprotocol/go-sdk/mcp"
1313
"github.com/shurcooL/githubv4"
1414
)
1515

@@ -309,8 +309,8 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
309309
contextTools := toolsets.NewToolset(ToolsetMetadataContext.ID, ToolsetMetadataContext.Description).
310310
AddReadTools(
311311
toolsets.NewServerTool(GetMe(getClient, t)),
312-
// toolsets.NewServerTool(GetTeams(getClient, getGQLClient, t)),
313-
// toolsets.NewServerTool(GetTeamMembers(getGQLClient, t)),
312+
toolsets.NewServerTool(GetTeams(getClient, getGQLClient, t)),
313+
toolsets.NewServerTool(GetTeamMembers(getGQLClient, t)),
314314
)
315315

316316
// gists := toolsets.NewToolset(ToolsetMetadataGists.ID, ToolsetMetadataGists.Description).
@@ -383,14 +383,14 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
383383
}
384384

385385
// InitDynamicToolset creates a dynamic toolset that can be used to enable other toolsets, and so requires the server and toolset group as arguments
386-
func InitDynamicToolset(s *server.MCPServer, tsg *toolsets.ToolsetGroup, t translations.TranslationHelperFunc) *toolsets.Toolset {
386+
func InitDynamicToolset(s *mcp.Server, tsg *toolsets.ToolsetGroup, t translations.TranslationHelperFunc) *toolsets.Toolset {
387387
// Create a new dynamic toolset
388388
// Need to add the dynamic toolset last so it can be used to enable other toolsets
389389
dynamicToolSelection := toolsets.NewToolset(ToolsetMetadataDynamic.ID, ToolsetMetadataDynamic.Description).
390390
AddReadTools(
391-
toolsets.NewServerTool(ListAvailableToolsets(tsg, t)),
392-
toolsets.NewServerTool(GetToolsetsTools(tsg, t)),
393-
toolsets.NewServerTool(EnableToolset(s, tsg, t)),
391+
// toolsets.NewServerTool(ListAvailableToolsets(tsg, t)),
392+
// toolsets.NewServerTool(GetToolsetsTools(tsg, t)),
393+
// toolsets.NewServerTool(EnableToolset(s, tsg, t)),
394394
)
395395

396396
dynamicToolSelection.Enabled = true

‎pkg/toolsets/toolsets.go‎

Lines changed: 44 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@ package toolsets
33
import (
44
"fmt"
55

6-
"github.com/mark3labs/mcp-go/mcp"
7-
"github.com/mark3labs/mcp-go/server"
6+
"github.com/modelcontextprotocol/go-sdk/mcp"
87
)
98

109
type ToolsetDoesNotExistError struct {
@@ -29,19 +28,34 @@ func NewToolsetDoesNotExistError(name string) *ToolsetDoesNotExistError {
2928
return &ToolsetDoesNotExistError{Name: name}
3029
}
3130

32-
func NewServerTool(tool mcp.Tool, handler server.ToolHandlerFunc) server.ServerTool {
33-
return server.ServerTool{Tool: tool, Handler: handler}
31+
type ServerTool struct {
32+
Tool mcp.Tool
33+
Handler mcp.ToolHandlerFor[map[string]any, any]
3434
}
3535

36-
func NewServerResourceTemplate(resourceTemplate mcp.ResourceTemplate, handler server.ResourceTemplateHandlerFunc) server.ServerResourceTemplate {
37-
return server.ServerResourceTemplate{
36+
func NewServerTool(tool mcp.Tool, handler mcp.ToolHandlerFor[map[string]any, any]) ServerTool {
37+
return ServerTool{Tool: tool, Handler: handler}
38+
}
39+
40+
type ServerResourceTemplate struct {
41+
Template mcp.ResourceTemplate
42+
Handler mcp.ResourceHandler
43+
}
44+
45+
func NewServerResourceTemplate(resourceTemplate mcp.ResourceTemplate, handler mcp.ResourceHandler) ServerResourceTemplate {
46+
return ServerResourceTemplate{
3847
Template: resourceTemplate,
3948
Handler: handler,
4049
}
4150
}
4251

43-
func NewServerPrompt(prompt mcp.Prompt, handler server.PromptHandlerFunc) server.ServerPrompt {
44-
return server.ServerPrompt{
52+
type ServerPrompt struct {
53+
Prompt mcp.Prompt
54+
Handler mcp.PromptHandler
55+
}
56+
57+
func NewServerPrompt(prompt mcp.Prompt, handler mcp.PromptHandler) ServerPrompt {
58+
return ServerPrompt{
4559
Prompt: prompt,
4660
Handler: handler,
4761
}
@@ -53,16 +67,16 @@ type Toolset struct {
5367
Description string
5468
Enabled bool
5569
readOnly bool
56-
writeTools []server.ServerTool
57-
readTools []server.ServerTool
70+
writeTools []ServerTool
71+
readTools []ServerTool
5872
// resources are not tools, but the community seems to be moving towards namespaces as a broader concept
5973
// and in order to have multiple servers running concurrently, we want to avoid overlapping resources too.
60-
resourceTemplates []server.ServerResourceTemplate
74+
resourceTemplates []ServerResourceTemplate
6175
// prompts are also not tools but are namespaced similarly
62-
prompts []server.ServerPrompt
76+
prompts []ServerPrompt
6377
}
6478

65-
func (t *Toolset) GetActiveTools() []server.ServerTool {
79+
func (t *Toolset) GetActiveTools() []ServerTool {
6680
if t.Enabled {
6781
if t.readOnly {
6882
return t.readTools
@@ -72,63 +86,63 @@ func (t *Toolset) GetActiveTools() []server.ServerTool {
7286
return nil
7387
}
7488

75-
func (t *Toolset) GetAvailableTools() []server.ServerTool {
89+
func (t *Toolset) GetAvailableTools() []ServerTool {
7690
if t.readOnly {
7791
return t.readTools
7892
}
7993
return append(t.readTools, t.writeTools...)
8094
}
8195

82-
func (t *Toolset) RegisterTools(s *server.MCPServer) {
96+
func (t *Toolset) RegisterTools(s *mcp.Server) {
8397
if !t.Enabled {
8498
return
8599
}
86100
for _, tool := range t.readTools {
87-
s.AddTool(tool.Tool, tool.Handler)
101+
mcp.AddTool(s, &tool.Tool, tool.Handler)
88102
}
89103
if !t.readOnly {
90104
for _, tool := range t.writeTools {
91-
s.AddTool(tool.Tool, tool.Handler)
105+
mcp.AddTool(s, &tool.Tool, tool.Handler)
92106
}
93107
}
94108
}
95109

96-
func (t *Toolset) AddResourceTemplates(templates ...server.ServerResourceTemplate) *Toolset {
110+
func (t *Toolset) AddResourceTemplates(templates ...ServerResourceTemplate) *Toolset {
97111
t.resourceTemplates = append(t.resourceTemplates, templates...)
98112
return t
99113
}
100114

101-
func (t *Toolset) AddPrompts(prompts ...server.ServerPrompt) *Toolset {
115+
func (t *Toolset) AddPrompts(prompts ...ServerPrompt) *Toolset {
102116
t.prompts = append(t.prompts, prompts...)
103117
return t
104118
}
105119

106-
func (t *Toolset) GetActiveResourceTemplates() []server.ServerResourceTemplate {
120+
func (t *Toolset) GetActiveResourceTemplates() []ServerResourceTemplate {
107121
if !t.Enabled {
108122
return nil
109123
}
110124
return t.resourceTemplates
111125
}
112126

113-
func (t *Toolset) GetAvailableResourceTemplates() []server.ServerResourceTemplate {
127+
func (t *Toolset) GetAvailableResourceTemplates() []ServerResourceTemplate {
114128
return t.resourceTemplates
115129
}
116130

117-
func (t *Toolset) RegisterResourcesTemplates(s *server.MCPServer) {
131+
func (t *Toolset) RegisterResourcesTemplates(s *mcp.Server) {
118132
if !t.Enabled {
119133
return
120134
}
121135
for _, resource := range t.resourceTemplates {
122-
s.AddResourceTemplate(resource.Template, resource.Handler)
136+
s.AddResourceTemplate(&resource.Template, resource.Handler)
123137
}
124138
}
125139

126-
func (t *Toolset) RegisterPrompts(s *server.MCPServer) {
140+
func (t *Toolset) RegisterPrompts(s *mcp.Server) {
127141
if !t.Enabled {
128142
return
129143
}
130144
for _, prompt := range t.prompts {
131-
s.AddPrompt(prompt.Prompt, prompt.Handler)
145+
s.AddPrompt(&prompt.Prompt, prompt.Handler)
132146
}
133147
}
134148

@@ -137,10 +151,10 @@ func (t *Toolset) SetReadOnly() {
137151
t.readOnly = true
138152
}
139153

140-
func (t *Toolset) AddWriteTools(tools ...server.ServerTool) *Toolset {
154+
func (t *Toolset) AddWriteTools(tools ...ServerTool) *Toolset {
141155
// Silently ignore if the toolset is read-only to avoid any breach of that contract
142156
for _, tool := range tools {
143-
if *tool.Tool.Annotations.ReadOnlyHint {
157+
if tool.Tool.Annotations.ReadOnlyHint {
144158
panic(fmt.Sprintf("tool (%s) is incorrectly annotated as read-only", tool.Tool.Name))
145159
}
146160
}
@@ -150,9 +164,9 @@ func (t *Toolset) AddWriteTools(tools ...server.ServerTool) *Toolset {
150164
return t
151165
}
152166

153-
func (t *Toolset) AddReadTools(tools ...server.ServerTool) *Toolset {
167+
func (t *Toolset) AddReadTools(tools ...ServerTool) *Toolset {
154168
for _, tool := range tools {
155-
if !*tool.Tool.Annotations.ReadOnlyHint {
169+
if !tool.Tool.Annotations.ReadOnlyHint {
156170
panic(fmt.Sprintf("tool (%s) must be annotated as read-only", tool.Tool.Name))
157171
}
158172
}
@@ -248,7 +262,7 @@ func (tg *ToolsetGroup) EnableToolset(name string) error {
248262
return nil
249263
}
250264

251-
func (tg *ToolsetGroup) RegisterAll(s *server.MCPServer) {
265+
func (tg *ToolsetGroup) RegisterAll(s *mcp.Server) {
252266
for _, toolset := range tg.Toolsets {
253267
toolset.RegisterTools(s)
254268
toolset.RegisterResourcesTemplates(s)

0 commit comments

Comments
 (0)