Skip to content

Commit 9227a95

Browse files
omgitsadsomgitsads
committed
Add schema caching to Toolsets.
1 parent 6462dee commit 9227a95

File tree

7 files changed

+129
-48
lines changed

7 files changed

+129
-48
lines changed

‎cmd/github-mcp-server/generate_docs.go‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/github/github-mcp-server/pkg/github"
1313
"github.com/github/github-mcp-server/pkg/lockdown"
1414
"github.com/github/github-mcp-server/pkg/raw"
15+
"github.com/github/github-mcp-server/pkg/schema"
1516
"github.com/github/github-mcp-server/pkg/toolsets"
1617
"github.com/github/github-mcp-server/pkg/translations"
1718
gogithub "github.com/google/go-github/v79/github"
@@ -65,9 +66,12 @@ func generateReadmeDocs(readmePath string) error {
6566
// Create translation helper
6667
t, _ := translations.TranslationHelper()
6768

69+
// Create a schema cache
70+
schemaCache := schema.NewCache()
71+
6872
// Create toolset group with mock clients
6973
repoAccessCache := lockdown.GetInstance(nil)
70-
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache)
74+
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache, schemaCache)
7175

7276
// Generate toolsets documentation
7377
toolsetsDoc := generateToolsetsDoc(tsg)
@@ -305,10 +309,12 @@ func generateRemoteToolsetsDoc() string {
305309
// Create translation helper
306310
t, _ := translations.TranslationHelper()
307311

312+
// Create a schema cache
313+
schemaCache := schema.NewCache()
314+
308315
// Create toolset group with mock clients
309316
repoAccessCache := lockdown.GetInstance(nil)
310-
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache)
311-
317+
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache, schemaCache)
312318
// Generate table header
313319
buf.WriteString("| Name | Description | API URL | 1-Click Install (VS Code) | Read-only Link | 1-Click Read-only Install (VS Code) |\n")
314320
buf.WriteString("|----------------|--------------------------------------------------|-------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n")

‎internal/ghmcp/server.go‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"github.com/github/github-mcp-server/pkg/lockdown"
1919
mcplog "github.com/github/github-mcp-server/pkg/log"
2020
"github.com/github/github-mcp-server/pkg/raw"
21+
"github.com/github/github-mcp-server/pkg/schema"
2122
"github.com/github/github-mcp-server/pkg/translations"
2223
gogithub "github.com/google/go-github/v79/github"
2324
"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -150,6 +151,8 @@ func NewMCPServer(cfg MCPServerConfig) (*mcp.Server, error) {
150151
ghServer.AddReceivingMiddleware(addGitHubAPIErrorToContext)
151152
ghServer.AddReceivingMiddleware(addUserAgentsMiddleware(cfg, restClient, gqlHTTPClient))
152153

154+
schemaCache := schema.NewCache()
155+
153156
// Create default toolsets
154157
tsg := github.DefaultToolsetGroup(
155158
cfg.ReadOnly,
@@ -160,6 +163,7 @@ func NewMCPServer(cfg MCPServerConfig) (*mcp.Server, error) {
160163
cfg.ContentWindowSize,
161164
github.FeatureFlags{LockdownMode: cfg.LockdownMode},
162165
repoAccessCache,
166+
schemaCache,
163167
)
164168

165169
// Enable and register toolsets if configured
@@ -180,15 +184,15 @@ func NewMCPServer(cfg MCPServerConfig) (*mcp.Server, error) {
180184
enabledTools := github.CleanTools(cfg.EnabledTools)
181185

182186
// Register the specified tools (additive to any toolsets already enabled)
183-
err = tsg.RegisterSpecificTools(ghServer, enabledTools, cfg.ReadOnly)
187+
err = tsg.RegisterSpecificTools(ghServer, enabledTools, cfg.ReadOnly, schemaCache)
184188
if err != nil {
185189
return nil, fmt.Errorf("failed to register tools: %w", err)
186190
}
187191
}
188192

189193
// Register dynamic toolsets if configured (additive to toolsets and tools)
190194
if cfg.DynamicToolsets {
191-
dynamic := github.InitDynamicToolset(ghServer, tsg, cfg.Translator)
195+
dynamic := github.InitDynamicToolset(ghServer, tsg, cfg.Translator, schemaCache)
192196
dynamic.RegisterTools(ghServer)
193197
}
194198

‎pkg/github/dynamic_tools.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"encoding/json"
66
"fmt"
77

8+
"github.com/github/github-mcp-server/pkg/schema"
89
"github.com/github/github-mcp-server/pkg/toolsets"
910
"github.com/github/github-mcp-server/pkg/translations"
1011
"github.com/github/github-mcp-server/pkg/utils"
@@ -20,7 +21,7 @@ func ToolsetEnum(toolsetGroup *toolsets.ToolsetGroup) []any {
2021
return toolsetNames
2122
}
2223

23-
func EnableToolset(s *mcp.Server, toolsetGroup *toolsets.ToolsetGroup, t translations.TranslationHelperFunc) (mcp.Tool, mcp.ToolHandlerFor[map[string]any, any]) {
24+
func EnableToolset(s *mcp.Server, toolsetGroup *toolsets.ToolsetGroup, t translations.TranslationHelperFunc, schemaCache *schema.Cache) (mcp.Tool, mcp.ToolHandlerFor[map[string]any, any]) {
2425
return mcp.Tool{
2526
Name: "enable_toolset",
2627
Description: t("TOOL_ENABLE_TOOLSET_DESCRIPTION", "Enable one of the sets of tools the GitHub MCP server provides, use get_toolset_tools and list_available_toolsets first to see what this will enable"),
@@ -62,7 +63,7 @@ func EnableToolset(s *mcp.Server, toolsetGroup *toolsets.ToolsetGroup, t transla
6263
// Send notification to all initialized sessions
6364
// s.sendNotificationToAllClients("notifications/tools/list_changed", nil)
6465
for _, serverTool := range toolset.GetActiveTools() {
65-
serverTool.RegisterFunc(s)
66+
serverTool.RegisterFunc(s, schemaCache)
6667
}
6768

6869
return utils.NewToolResultText(fmt.Sprintf("Toolset %s enabled", toolsetName)), nil, nil

‎pkg/github/tools.go‎

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77

88
"github.com/github/github-mcp-server/pkg/lockdown"
99
"github.com/github/github-mcp-server/pkg/raw"
10+
"github.com/github/github-mcp-server/pkg/schema"
1011
"github.com/github/github-mcp-server/pkg/toolsets"
1112
"github.com/github/github-mcp-server/pkg/translations"
1213
"github.com/google/go-github/v79/github"
@@ -160,12 +161,21 @@ func GetDefaultToolsetIDs() []string {
160161
}
161162
}
162163

163-
func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetGQLClientFn, getRawClient raw.GetRawClientFn, t translations.TranslationHelperFunc, contentWindowSize int, flags FeatureFlags, cache *lockdown.RepoAccessCache) *toolsets.ToolsetGroup {
164+
func DefaultToolsetGroup(readOnly bool,
165+
getClient GetClientFn,
166+
getGQLClient GetGQLClientFn,
167+
getRawClient raw.GetRawClientFn,
168+
t translations.TranslationHelperFunc,
169+
contentWindowSize int,
170+
flags FeatureFlags,
171+
cache *lockdown.RepoAccessCache,
172+
schemaCache *schema.Cache,
173+
) *toolsets.ToolsetGroup {
164174
tsg := toolsets.NewToolsetGroup(readOnly)
165175

166176
// Define all available features with their default state (disabled)
167177
// Create toolsets
168-
repos := toolsets.NewToolset(ToolsetMetadataRepos.ID, ToolsetMetadataRepos.Description).
178+
repos := toolsets.NewToolset(ToolsetMetadataRepos.ID, ToolsetMetadataRepos.Description, schemaCache).
169179
AddReadTools(
170180
toolsets.NewServerTool(SearchRepositories(getClient, t)),
171181
toolsets.NewServerTool(GetFileContents(getClient, getRawClient, t)),
@@ -194,11 +204,11 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
194204
toolsets.NewServerResourceTemplate(GetRepositoryResourceTagContent(getClient, getRawClient, t)),
195205
toolsets.NewServerResourceTemplate(GetRepositoryResourcePrContent(getClient, getRawClient, t)),
196206
)
197-
git := toolsets.NewToolset(ToolsetMetadataGit.ID, ToolsetMetadataGit.Description).
207+
git := toolsets.NewToolset(ToolsetMetadataGit.ID, ToolsetMetadataGit.Description, schemaCache).
198208
AddReadTools(
199209
toolsets.NewServerTool(GetRepositoryTree(getClient, t)),
200210
)
201-
issues := toolsets.NewToolset(ToolsetMetadataIssues.ID, ToolsetMetadataIssues.Description).
211+
issues := toolsets.NewToolset(ToolsetMetadataIssues.ID, ToolsetMetadataIssues.Description, schemaCache).
202212
AddReadTools(
203213
toolsets.NewServerTool(IssueRead(getClient, getGQLClient, cache, t, flags)),
204214
toolsets.NewServerTool(SearchIssues(getClient, t)),
@@ -215,15 +225,15 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
215225
toolsets.NewServerPrompt(AssignCodingAgentPrompt(t)),
216226
toolsets.NewServerPrompt(IssueToFixWorkflowPrompt(t)),
217227
)
218-
users := toolsets.NewToolset(ToolsetMetadataUsers.ID, ToolsetMetadataUsers.Description).
228+
users := toolsets.NewToolset(ToolsetMetadataUsers.ID, ToolsetMetadataUsers.Description, schemaCache).
219229
AddReadTools(
220230
toolsets.NewServerTool(SearchUsers(getClient, t)),
221231
)
222-
orgs := toolsets.NewToolset(ToolsetMetadataOrgs.ID, ToolsetMetadataOrgs.Description).
232+
orgs := toolsets.NewToolset(ToolsetMetadataOrgs.ID, ToolsetMetadataOrgs.Description, schemaCache).
223233
AddReadTools(
224234
toolsets.NewServerTool(SearchOrgs(getClient, t)),
225235
)
226-
pullRequests := toolsets.NewToolset(ToolsetMetadataPullRequests.ID, ToolsetMetadataPullRequests.Description).
236+
pullRequests := toolsets.NewToolset(ToolsetMetadataPullRequests.ID, ToolsetMetadataPullRequests.Description, schemaCache).
227237
AddReadTools(
228238
toolsets.NewServerTool(PullRequestRead(getClient, cache, t, flags)),
229239
toolsets.NewServerTool(ListPullRequests(getClient, t)),
@@ -239,23 +249,23 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
239249
toolsets.NewServerTool(PullRequestReviewWrite(getGQLClient, t)),
240250
toolsets.NewServerTool(AddCommentToPendingReview(getGQLClient, t)),
241251
)
242-
codeSecurity := toolsets.NewToolset(ToolsetMetadataCodeSecurity.ID, ToolsetMetadataCodeSecurity.Description).
252+
codeSecurity := toolsets.NewToolset(ToolsetMetadataCodeSecurity.ID, ToolsetMetadataCodeSecurity.Description, schemaCache).
243253
AddReadTools(
244254
toolsets.NewServerTool(GetCodeScanningAlert(getClient, t)),
245255
toolsets.NewServerTool(ListCodeScanningAlerts(getClient, t)),
246256
)
247-
secretProtection := toolsets.NewToolset(ToolsetMetadataSecretProtection.ID, ToolsetMetadataSecretProtection.Description).
257+
secretProtection := toolsets.NewToolset(ToolsetMetadataSecretProtection.ID, ToolsetMetadataSecretProtection.Description, schemaCache).
248258
AddReadTools(
249259
toolsets.NewServerTool(GetSecretScanningAlert(getClient, t)),
250260
toolsets.NewServerTool(ListSecretScanningAlerts(getClient, t)),
251261
)
252-
dependabot := toolsets.NewToolset(ToolsetMetadataDependabot.ID, ToolsetMetadataDependabot.Description).
262+
dependabot := toolsets.NewToolset(ToolsetMetadataDependabot.ID, ToolsetMetadataDependabot.Description, schemaCache).
253263
AddReadTools(
254264
toolsets.NewServerTool(GetDependabotAlert(getClient, t)),
255265
toolsets.NewServerTool(ListDependabotAlerts(getClient, t)),
256266
)
257267

258-
notifications := toolsets.NewToolset(ToolsetMetadataNotifications.ID, ToolsetMetadataNotifications.Description).
268+
notifications := toolsets.NewToolset(ToolsetMetadataNotifications.ID, ToolsetMetadataNotifications.Description, schemaCache).
259269
AddReadTools(
260270
toolsets.NewServerTool(ListNotifications(getClient, t)),
261271
toolsets.NewServerTool(GetNotificationDetails(getClient, t)),
@@ -267,15 +277,15 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
267277
toolsets.NewServerTool(ManageRepositoryNotificationSubscription(getClient, t)),
268278
)
269279

270-
discussions := toolsets.NewToolset(ToolsetMetadataDiscussions.ID, ToolsetMetadataDiscussions.Description).
280+
discussions := toolsets.NewToolset(ToolsetMetadataDiscussions.ID, ToolsetMetadataDiscussions.Description, schemaCache).
271281
AddReadTools(
272282
toolsets.NewServerTool(ListDiscussions(getGQLClient, t)),
273283
toolsets.NewServerTool(GetDiscussion(getGQLClient, t)),
274284
toolsets.NewServerTool(GetDiscussionComments(getGQLClient, t)),
275285
toolsets.NewServerTool(ListDiscussionCategories(getGQLClient, t)),
276286
)
277287

278-
actions := toolsets.NewToolset(ToolsetMetadataActions.ID, ToolsetMetadataActions.Description).
288+
actions := toolsets.NewToolset(ToolsetMetadataActions.ID, ToolsetMetadataActions.Description, schemaCache).
279289
AddReadTools(
280290
toolsets.NewServerTool(ListWorkflows(getClient, t)),
281291
toolsets.NewServerTool(ListWorkflowRuns(getClient, t)),
@@ -295,7 +305,7 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
295305
toolsets.NewServerTool(DeleteWorkflowRunLogs(getClient, t)),
296306
)
297307

298-
securityAdvisories := toolsets.NewToolset(ToolsetMetadataSecurityAdvisories.ID, ToolsetMetadataSecurityAdvisories.Description).
308+
securityAdvisories := toolsets.NewToolset(ToolsetMetadataSecurityAdvisories.ID, ToolsetMetadataSecurityAdvisories.Description, schemaCache).
299309
AddReadTools(
300310
toolsets.NewServerTool(ListGlobalSecurityAdvisories(getClient, t)),
301311
toolsets.NewServerTool(GetGlobalSecurityAdvisory(getClient, t)),
@@ -304,16 +314,16 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
304314
)
305315

306316
// // Keep experiments alive so the system doesn't error out when it's always enabled
307-
experiments := toolsets.NewToolset(ToolsetMetadataExperiments.ID, ToolsetMetadataExperiments.Description)
317+
experiments := toolsets.NewToolset(ToolsetMetadataExperiments.ID, ToolsetMetadataExperiments.Description, schemaCache)
308318

309-
contextTools := toolsets.NewToolset(ToolsetMetadataContext.ID, ToolsetMetadataContext.Description).
319+
contextTools := toolsets.NewToolset(ToolsetMetadataContext.ID, ToolsetMetadataContext.Description, schemaCache).
310320
AddReadTools(
311321
toolsets.NewServerTool(GetMe(getClient, t)),
312322
toolsets.NewServerTool(GetTeams(getClient, getGQLClient, t)),
313323
toolsets.NewServerTool(GetTeamMembers(getGQLClient, t)),
314324
)
315325

316-
gists := toolsets.NewToolset(ToolsetMetadataGists.ID, ToolsetMetadataGists.Description).
326+
gists := toolsets.NewToolset(ToolsetMetadataGists.ID, ToolsetMetadataGists.Description, schemaCache).
317327
AddReadTools(
318328
toolsets.NewServerTool(ListGists(getClient, t)),
319329
toolsets.NewServerTool(GetGist(getClient, t)),
@@ -323,7 +333,7 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
323333
toolsets.NewServerTool(UpdateGist(getClient, t)),
324334
)
325335

326-
projects := toolsets.NewToolset(ToolsetMetadataProjects.ID, ToolsetMetadataProjects.Description).
336+
projects := toolsets.NewToolset(ToolsetMetadataProjects.ID, ToolsetMetadataProjects.Description, schemaCache).
327337
AddReadTools(
328338
toolsets.NewServerTool(ListProjects(getClient, t)),
329339
toolsets.NewServerTool(GetProject(getClient, t)),
@@ -337,15 +347,15 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
337347
toolsets.NewServerTool(DeleteProjectItem(getClient, t)),
338348
toolsets.NewServerTool(UpdateProjectItem(getClient, t)),
339349
)
340-
stargazers := toolsets.NewToolset(ToolsetMetadataStargazers.ID, ToolsetMetadataStargazers.Description).
350+
stargazers := toolsets.NewToolset(ToolsetMetadataStargazers.ID, ToolsetMetadataStargazers.Description, schemaCache).
341351
AddReadTools(
342352
toolsets.NewServerTool(ListStarredRepositories(getClient, t)),
343353
).
344354
AddWriteTools(
345355
toolsets.NewServerTool(StarRepository(getClient, t)),
346356
toolsets.NewServerTool(UnstarRepository(getClient, t)),
347357
)
348-
labels := toolsets.NewToolset(ToolsetLabels.ID, ToolsetLabels.Description).
358+
labels := toolsets.NewToolset(ToolsetLabels.ID, ToolsetLabels.Description, schemaCache).
349359
AddReadTools(
350360
// get
351361
toolsets.NewServerTool(GetLabel(getGQLClient, t)),
@@ -384,14 +394,14 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
384394
// InitDynamicToolset creates a dynamic toolset that can be used to enable other toolsets, and so requires the server and toolset group as arguments
385395
//
386396
//nolint:unused
387-
func InitDynamicToolset(s *mcp.Server, tsg *toolsets.ToolsetGroup, t translations.TranslationHelperFunc) *toolsets.Toolset {
397+
func InitDynamicToolset(s *mcp.Server, tsg *toolsets.ToolsetGroup, t translations.TranslationHelperFunc, schemaCache *schema.Cache) *toolsets.Toolset {
388398
// Create a new dynamic toolset
389399
// Need to add the dynamic toolset last so it can be used to enable other toolsets
390-
dynamicToolSelection := toolsets.NewToolset(ToolsetMetadataDynamic.ID, ToolsetMetadataDynamic.Description).
400+
dynamicToolSelection := toolsets.NewToolset(ToolsetMetadataDynamic.ID, ToolsetMetadataDynamic.Description, schemaCache).
391401
AddReadTools(
392402
toolsets.NewServerTool(ListAvailableToolsets(tsg, t)),
393403
toolsets.NewServerTool(GetToolsetsTools(tsg, t)),
394-
toolsets.NewServerTool(EnableToolset(s, tsg, t)),
404+
toolsets.NewServerTool(EnableToolset(s, tsg, t, schemaCache)),
395405
)
396406

397407
dynamicToolSelection.Enabled = true

‎pkg/schema/cache.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package schema
2+
3+
import (
4+
"sync"
5+
6+
"github.com/google/jsonschema-go/jsonschema"
7+
)
8+
9+
type Cache struct {
10+
cachedSchemas map[string]*jsonschema.Resolved
11+
mu sync.Mutex
12+
}
13+
14+
func NewCache() *Cache {
15+
return &Cache{
16+
cachedSchemas: make(map[string]*jsonschema.Resolved),
17+
}
18+
}
19+
20+
func (sc *Cache) GetCachedOrResolveSchema(toolName string, schema *jsonschema.Schema) (*jsonschema.Resolved, error) {
21+
if resolvedSchema, ok := sc.cachedSchemas[toolName]; ok {
22+
return resolvedSchema, nil
23+
}
24+
25+
sc.mu.Lock()
26+
defer sc.mu.Unlock()
27+
resolvedSchema, err := schema.Resolve(nil)
28+
if err != nil {
29+
return nil, err
30+
}
31+
32+
sc.cachedSchemas[toolName] = resolvedSchema
33+
return resolvedSchema, nil
34+
}

0 commit comments

Comments
 (0)