Skip to content

Commit 5562335

Browse files
JoannaaKLJoannaaKL
committed
Add cache
1 parent 2f7f3ab commit 5562335

File tree

11 files changed

+496
-91
lines changed

11 files changed

+496
-91
lines changed

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"strings"
1111

1212
"github.com/github/github-mcp-server/pkg/github"
13+
"github.com/github/github-mcp-server/pkg/lockdown"
1314
"github.com/github/github-mcp-server/pkg/raw"
1415
"github.com/github/github-mcp-server/pkg/toolsets"
1516
"github.com/github/github-mcp-server/pkg/translations"
@@ -64,7 +65,8 @@ func generateReadmeDocs(readmePath string) error {
6465
t, _ := translations.TranslationHelper()
6566

6667
// Create toolset group with mock clients
67-
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{})
68+
repoAccessCache := lockdown.NewRepoAccessCache(nil)
69+
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache)
6870

6971
// Generate toolsets documentation
7072
toolsetsDoc := generateToolsetsDoc(tsg)
@@ -302,7 +304,8 @@ func generateRemoteToolsetsDoc() string {
302304
t, _ := translations.TranslationHelper()
303305

304306
// Create toolset group with mock clients
305-
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{})
307+
repoAccessCache := lockdown.NewRepoAccessCache(nil)
308+
tsg := github.DefaultToolsetGroup(false, mockGetClient, mockGetGQLClient, mockGetRawClient, t, 5000, github.FeatureFlags{}, repoAccessCache)
306309

307310
// Generate table header
308311
buf.WriteString("| Name | Description | API URL | 1-Click Install (VS Code) | Read-only Link | 1-Click Read-only Install (VS Code) |\n")

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"os"
77
"strings"
8+
"time"
89

910
"github.com/github/github-mcp-server/internal/ghmcp"
1011
"github.com/github/github-mcp-server/pkg/github"
@@ -50,6 +51,7 @@ var (
5051
enabledToolsets = []string{github.ToolsetMetadataDefault.ID}
5152
}
5253

54+
ttl := viper.GetDuration("repo-access-cache-ttl")
5355
stdioServerConfig := ghmcp.StdioServerConfig{
5456
Version: version,
5557
Host: viper.GetString("host"),
@@ -62,6 +64,7 @@ var (
6264
LogFilePath: viper.GetString("log-file"),
6365
ContentWindowSize: viper.GetInt("content-window-size"),
6466
LockdownMode: viper.GetBool("lockdown-mode"),
67+
RepoAccessCacheTTL: &ttl,
6568
}
6669
return ghmcp.RunStdioServer(stdioServerConfig)
6770
},
@@ -84,6 +87,7 @@ func init() {
8487
rootCmd.PersistentFlags().String("gh-host", "", "Specify the GitHub hostname (for GitHub Enterprise etc.)")
8588
rootCmd.PersistentFlags().Int("content-window-size", 5000, "Specify the content window size")
8689
rootCmd.PersistentFlags().Bool("lockdown-mode", false, "Enable lockdown mode")
90+
rootCmd.PersistentFlags().Duration("repo-access-cache-ttl", 5*time.Minute, "Override the repo access cache TTL (e.g. 1m, 0s to disable)")
8791

8892
// Bind flag to viper
8993
_ = viper.BindPFlag("toolsets", rootCmd.PersistentFlags().Lookup("toolsets"))
@@ -95,6 +99,7 @@ func init() {
9599
_ = viper.BindPFlag("host", rootCmd.PersistentFlags().Lookup("gh-host"))
96100
_ = viper.BindPFlag("content-window-size", rootCmd.PersistentFlags().Lookup("content-window-size"))
97101
_ = viper.BindPFlag("lockdown-mode", rootCmd.PersistentFlags().Lookup("lockdown-mode"))
102+
_ = viper.BindPFlag("repo-access-cache-ttl", rootCmd.PersistentFlags().Lookup("repo-access-cache-ttl"))
98103

99104
// Add subcommands
100105
rootCmd.AddCommand(stdioCmd)

‎internal/ghmcp/server.go‎

Lines changed: 34 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616

1717
"github.com/github/github-mcp-server/pkg/errors"
1818
"github.com/github/github-mcp-server/pkg/github"
19+
"github.com/github/github-mcp-server/pkg/lockdown"
1920
mcplog "github.com/github/github-mcp-server/pkg/log"
2021
"github.com/github/github-mcp-server/pkg/raw"
2122
"github.com/github/github-mcp-server/pkg/translations"
@@ -54,6 +55,9 @@ type MCPServerConfig struct {
5455

5556
// LockdownMode indicates if we should enable lockdown mode
5657
LockdownMode bool
58+
59+
// RepoAccessTTL overrides the default TTL for repository access cache entries.
60+
RepoAccessTTL *time.Duration
5761
}
5862

5963
const stdioServerLogPrefix = "stdioserver"
@@ -80,6 +84,14 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
8084
},
8185
} // We're going to wrap the Transport later in beforeInit
8286
gqlClient := githubv4.NewEnterpriseClient(apiHost.graphqlURL.String(), gqlHTTPClient)
87+
repoAccessOpts := []lockdown.RepoAccessOption{}
88+
if cfg.RepoAccessTTL != nil {
89+
repoAccessOpts = append(repoAccessOpts, lockdown.WithTTL(*cfg.RepoAccessTTL))
90+
}
91+
var repoAccessCache *lockdown.RepoAccessCache
92+
if cfg.LockdownMode {
93+
repoAccessCache = lockdown.NewRepoAccessCache(gqlClient, repoAccessOpts...)
94+
}
8395

8496
// When a client send an initialize request, update the user agent to include the client info.
8597
beforeInit := func(_ context.Context, _ any, message *mcp.InitializeRequest) {
@@ -165,6 +177,7 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
165177
cfg.Translator,
166178
cfg.ContentWindowSize,
167179
github.FeatureFlags{LockdownMode: cfg.LockdownMode},
180+
repoAccessCache,
168181
)
169182
err = tsg.EnableToolsets(enabledToolsets, nil)
170183

@@ -219,6 +232,9 @@ type StdioServerConfig struct {
219232

220233
// LockdownMode indicates if we should enable lockdown mode
221234
LockdownMode bool
235+
236+
// RepoAccessCacheTTL overrides the default TTL for repository access cache entries.
237+
RepoAccessCacheTTL *time.Duration
222238
}
223239

224240
// RunStdioServer is not concurrent safe.
@@ -229,23 +245,6 @@ func RunStdioServer(cfg StdioServerConfig) error {
229245

230246
t, dumpTranslations := translations.TranslationHelper()
231247

232-
ghServer, err := NewMCPServer(MCPServerConfig{
233-
Version: cfg.Version,
234-
Host: cfg.Host,
235-
Token: cfg.Token,
236-
EnabledToolsets: cfg.EnabledToolsets,
237-
DynamicToolsets: cfg.DynamicToolsets,
238-
ReadOnly: cfg.ReadOnly,
239-
Translator: t,
240-
ContentWindowSize: cfg.ContentWindowSize,
241-
LockdownMode: cfg.LockdownMode,
242-
})
243-
if err != nil {
244-
return fmt.Errorf("failed to create MCP server: %w", err)
245-
}
246-
247-
stdioServer := server.NewStdioServer(ghServer)
248-
249248
var slogHandler slog.Handler
250249
var logOutput io.Writer
251250
if cfg.LogFilePath != "" {
@@ -262,6 +261,24 @@ func RunStdioServer(cfg StdioServerConfig) error {
262261
logger := slog.New(slogHandler)
263262
logger.Info("starting server", "version", cfg.Version, "host", cfg.Host, "dynamicToolsets", cfg.DynamicToolsets, "readOnly", cfg.ReadOnly, "lockdownEnabled", cfg.LockdownMode)
264263
stdLogger := log.New(logOutput, stdioServerLogPrefix, 0)
264+
265+
ghServer, err := NewMCPServer(MCPServerConfig{
266+
Version: cfg.Version,
267+
Host: cfg.Host,
268+
Token: cfg.Token,
269+
EnabledToolsets: cfg.EnabledToolsets,
270+
DynamicToolsets: cfg.DynamicToolsets,
271+
ReadOnly: cfg.ReadOnly,
272+
Translator: t,
273+
ContentWindowSize: cfg.ContentWindowSize,
274+
LockdownMode: cfg.LockdownMode,
275+
RepoAccessTTL: cfg.RepoAccessCacheTTL,
276+
})
277+
if err != nil {
278+
return fmt.Errorf("failed to create MCP server: %w", err)
279+
}
280+
281+
stdioServer := server.NewStdioServer(ghServer)
265282
stdioServer.SetErrorLogger(stdLogger)
266283

267284
if cfg.ExportTranslations {

‎pkg/github/issues.go‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,7 @@ func fragmentToIssue(fragment IssueFragment) *github.Issue {
228228
}
229229

230230
// GetIssue creates a tool to get details of a specific issue in a GitHub repository.
231-
func IssueRead(getClient GetClientFn, getGQLClient GetGQLClientFn, t translations.TranslationHelperFunc, flags FeatureFlags) (tool mcp.Tool, handler server.ToolHandlerFunc) {
231+
func IssueRead(getClient GetClientFn, getGQLClient GetGQLClientFn, cache *lockdown.RepoAccessCache, t translations.TranslationHelperFunc, flags FeatureFlags) (tool mcp.Tool, handler server.ToolHandlerFunc) {
232232
return mcp.NewTool("issue_read",
233233
mcp.WithDescription(t("TOOL_ISSUE_READ_DESCRIPTION", "Get information about a specific issue in a GitHub repository.")),
234234
mcp.WithToolAnnotation(mcp.ToolAnnotation{
@@ -297,11 +297,11 @@ Options are:
297297

298298
switch method {
299299
case "get":
300-
return GetIssue(ctx, client, gqlClient, owner, repo, issueNumber, flags)
300+
return GetIssue(ctx, client, cache, owner, repo, issueNumber, flags)
301301
case "get_comments":
302-
return GetIssueComments(ctx, client, gqlClient, owner, repo, issueNumber, pagination, flags)
302+
return GetIssueComments(ctx, client, cache, owner, repo, issueNumber, pagination, flags)
303303
case "get_sub_issues":
304-
return GetSubIssues(ctx, client, gqlClient, owner, repo, issueNumber, pagination, flags)
304+
return GetSubIssues(ctx, client, cache, owner, repo, issueNumber, pagination, flags)
305305
case "get_labels":
306306
return GetIssueLabels(ctx, gqlClient, owner, repo, issueNumber)
307307
default:
@@ -310,7 +310,7 @@ Options are:
310310
}
311311
}
312312

313-
func GetIssue(ctx context.Context, client *github.Client, gqlClient *githubv4.Client, owner string, repo string, issueNumber int, flags FeatureFlags) (*mcp.CallToolResult, error) {
313+
func GetIssue(ctx context.Context, client *github.Client, cache *lockdown.RepoAccessCache, owner string, repo string, issueNumber int, flags FeatureFlags) (*mcp.CallToolResult, error) {
314314
issue, resp, err := client.Issues.Get(ctx, owner, repo, issueNumber)
315315
if err != nil {
316316
return nil, fmt.Errorf("failed to get issue: %w", err)
@@ -326,8 +326,12 @@ func GetIssue(ctx context.Context, client *github.Client, gqlClient *githubv4.Cl
326326
}
327327

328328
if flags.LockdownMode {
329-
if issue.User != nil {
330-
isPrivate, hasPushAccess, err := lockdown.GetRepoAccessInfo(ctx, gqlClient, *issue.User.Login, owner, repo)
329+
if cache == nil {
330+
return nil, fmt.Errorf("lockdown cache is not configured")
331+
}
332+
login := issue.GetUser().GetLogin()
333+
if login != "" {
334+
isPrivate, hasPushAccess, err := cache.GetRepoAccessInfo(ctx, login, owner, repo)
331335
if err != nil {
332336
return mcp.NewToolResultError(fmt.Sprintf("failed to check lockdown mode: %v", err)), nil
333337
}
@@ -355,7 +359,7 @@ func GetIssue(ctx context.Context, client *github.Client, gqlClient *githubv4.Cl
355359
return mcp.NewToolResultText(string(r)), nil
356360
}
357361

358-
func GetIssueComments(ctx context.Context, client *github.Client, gqlClient *githubv4.Client, owner string, repo string, issueNumber int, pagination PaginationParams, flags FeatureFlags) (*mcp.CallToolResult, error) {
362+
func GetIssueComments(ctx context.Context, client *github.Client, cache *lockdown.RepoAccessCache, owner string, repo string, issueNumber int, pagination PaginationParams, flags FeatureFlags) (*mcp.CallToolResult, error) {
359363
opts := &github.IssueListCommentsOptions{
360364
ListOptions: github.ListOptions{
361365
Page: pagination.Page,
@@ -377,9 +381,20 @@ func GetIssueComments(ctx context.Context, client *github.Client, gqlClient *git
377381
return mcp.NewToolResultError(fmt.Sprintf("failed to get issue comments: %s", string(body))), nil
378382
}
379383
if flags.LockdownMode {
380-
filteredComments := []*github.IssueComment{}
384+
if cache == nil {
385+
return nil, fmt.Errorf("lockdown cache is not configured")
386+
}
387+
filteredComments := make([]*github.IssueComment, 0, len(comments))
381388
for _, comment := range comments {
382-
isPrivate, hasPushAccess, err := lockdown.GetRepoAccessInfo(ctx, gqlClient, *comment.User.Login, owner, repo)
389+
user := comment.User
390+
if user == nil {
391+
continue
392+
}
393+
login := user.GetLogin()
394+
if login == "" {
395+
continue
396+
}
397+
isPrivate, hasPushAccess, err := cache.GetRepoAccessInfo(ctx, login, owner, repo)
383398
if err != nil {
384399
return mcp.NewToolResultError(fmt.Sprintf("failed to check lockdown mode: %v", err)), nil
385400
}
@@ -402,7 +417,7 @@ func GetIssueComments(ctx context.Context, client *github.Client, gqlClient *git
402417
return mcp.NewToolResultText(string(r)), nil
403418
}
404419

405-
func GetSubIssues(ctx context.Context, client *github.Client, gqlClient *githubv4.Client, owner string, repo string, issueNumber int, pagination PaginationParams, featureFlags FeatureFlags) (*mcp.CallToolResult, error) {
420+
func GetSubIssues(ctx context.Context, client *github.Client, cache *lockdown.RepoAccessCache, owner string, repo string, issueNumber int, pagination PaginationParams, featureFlags FeatureFlags) (*mcp.CallToolResult, error) {
406421
opts := &github.IssueListOptions{
407422
ListOptions: github.ListOptions{
408423
Page: pagination.Page,
@@ -430,9 +445,20 @@ func GetSubIssues(ctx context.Context, client *github.Client, gqlClient *githubv
430445
}
431446

432447
if featureFlags.LockdownMode {
433-
filteredSubIssues := []*github.SubIssue{}
448+
if cache == nil {
449+
return nil, fmt.Errorf("lockdown cache is not configured")
450+
}
451+
filteredSubIssues := make([]*github.SubIssue, 0, len(subIssues))
434452
for _, subIssue := range subIssues {
435-
isPrivate, hasPushAccess, err := lockdown.GetRepoAccessInfo(ctx, gqlClient, *subIssue.User.Login, owner, repo)
453+
user := subIssue.User
454+
if user == nil {
455+
continue
456+
}
457+
login := user.GetLogin()
458+
if login == "" {
459+
continue
460+
}
461+
isPrivate, hasPushAccess, err := cache.GetRepoAccessInfo(ctx, login, owner, repo)
436462
if err != nil {
437463
return mcp.NewToolResultError(fmt.Sprintf("failed to check lockdown mode: %v", err)), nil
438464
}

‎pkg/github/issues_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ func Test_GetIssue(t *testing.T) {
2323
// Verify tool definition once
2424
mockClient := github.NewClient(nil)
2525
defaultGQLClient := githubv4.NewClient(nil)
26-
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(defaultGQLClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
26+
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(defaultGQLClient), stubRepoAccessCache(defaultGQLClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
2727
require.NoError(t, toolsnaps.Test(tool.Name, tool))
2828

2929
assert.Equal(t, "issue_read", tool.Name)
@@ -212,7 +212,7 @@ func Test_GetIssue(t *testing.T) {
212212
}
213213

214214
flags := stubFeatureFlags(map[string]bool{"lockdown-mode": tc.lockdownEnabled})
215-
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, flags)
215+
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, flags)
216216

217217
request := createMCPRequest(tc.requestArgs)
218218
result, err := handler(context.Background(), request)
@@ -1710,7 +1710,7 @@ func Test_GetIssueComments(t *testing.T) {
17101710
// Verify tool definition once
17111711
mockClient := github.NewClient(nil)
17121712
gqlClient := githubv4.NewClient(nil)
1713-
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
1713+
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
17141714
require.NoError(t, toolsnaps.Test(tool.Name, tool))
17151715

17161716
assert.Equal(t, "issue_read", tool.Name)
@@ -1816,7 +1816,7 @@ func Test_GetIssueComments(t *testing.T) {
18161816
// Setup client with mock
18171817
client := github.NewClient(tc.mockedClient)
18181818
gqlClient := githubv4.NewClient(nil)
1819-
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
1819+
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
18201820

18211821
// Create call request
18221822
request := createMCPRequest(tc.requestArgs)
@@ -1853,7 +1853,7 @@ func Test_GetIssueLabels(t *testing.T) {
18531853
// Verify tool definition
18541854
mockGQClient := githubv4.NewClient(nil)
18551855
mockClient := github.NewClient(nil)
1856-
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(mockGQClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
1856+
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(mockGQClient), stubRepoAccessCache(mockGQClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
18571857
require.NoError(t, toolsnaps.Test(tool.Name, tool))
18581858

18591859
assert.Equal(t, "issue_read", tool.Name)
@@ -1928,7 +1928,7 @@ func Test_GetIssueLabels(t *testing.T) {
19281928
t.Run(tc.name, func(t *testing.T) {
19291929
gqlClient := githubv4.NewClient(tc.mockedClient)
19301930
client := github.NewClient(nil)
1931-
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
1931+
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
19321932

19331933
request := createMCPRequest(tc.requestArgs)
19341934
result, err := handler(context.Background(), request)
@@ -2619,7 +2619,7 @@ func Test_GetSubIssues(t *testing.T) {
26192619
// Verify tool definition once
26202620
mockClient := github.NewClient(nil)
26212621
gqlClient := githubv4.NewClient(nil)
2622-
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
2622+
tool, _ := IssueRead(stubGetClientFn(mockClient), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
26232623
require.NoError(t, toolsnaps.Test(tool.Name, tool))
26242624

26252625
assert.Equal(t, "issue_read", tool.Name)
@@ -2816,7 +2816,7 @@ func Test_GetSubIssues(t *testing.T) {
28162816
// Setup client with mock
28172817
client := github.NewClient(tc.mockedClient)
28182818
gqlClient := githubv4.NewClient(nil)
2819-
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
2819+
_, handler := IssueRead(stubGetClientFn(client), stubGetGQLClientFn(gqlClient), stubRepoAccessCache(gqlClient), translations.NullTranslationHelper, stubFeatureFlags(map[string]bool{"lockdown-mode": false}))
28202820

28212821
// Create call request
28222822
request := createMCPRequest(tc.requestArgs)

0 commit comments

Comments
 (0)