use GraphQL API to get file SHA

Author: LuluBeatson

pkg/github/repositories.go

@@ -16,8 +16,13 @@ import (
 	"github.com/google/go-github/v72/github"
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
+	"github.com/shurcooL/githubv4"
 )
 
+// getFileSHAFunc is a package-level variable that holds the getFileSHA function
+// This allows tests to mock the behavior
+var getFileSHAFunc = getFileSHA
+
 func GetCommit(getClient GetClientFn, t translations.TranslationHelperFunc) (tool mcp.Tool, handler server.ToolHandlerFunc) {
 	return mcp.NewTool("get_commit",
 			mcp.WithDescription(t("TOOL_GET_COMMITS_DESCRIPTION", "Get details for a commit from a GitHub repository")),
@@ -446,7 +451,7 @@ func CreateRepository(getClient GetClientFn, t translations.TranslationHelperFun
 }
 
 // GetFileContents creates a tool to get the contents of a file or directory from a GitHub repository.
-func GetFileContents(getClient GetClientFn, getRawClient raw.GetRawClientFn, t translations.TranslationHelperFunc) (tool mcp.Tool, handler server.ToolHandlerFunc) {
+func GetFileContents(getClient GetClientFn, getRawClient raw.GetRawClientFn, getGQLClient GetGQLClientFn, t translations.TranslationHelperFunc) (tool mcp.Tool, handler server.ToolHandlerFunc) {
 	return mcp.NewTool("get_file_contents",
 			mcp.WithDescription(t("TOOL_GET_FILE_CONTENTS_DESCRIPTION", "Get the contents of a file or directory from a GitHub repository")),
 			mcp.WithToolAnnotation(mcp.ToolAnnotation{
@@ -507,15 +512,13 @@ func GetFileContents(getClient GetClientFn, getRawClient raw.GetRawClientFn, t t
 			// If the path is (most likely) not to be a directory, we will
 			// first try to get the raw content from the GitHub raw content API.
 			if path != "" && !strings.HasSuffix(path, "/") {
-				// First, get file info from Contents API to retrieve SHA
-				var fileSHA string
-				opts := &github.RepositoryContentGetOptions{Ref: ref}
-				fileContent, _, respContents, err := client.Repositories.GetContents(ctx, owner, repo, path, opts)
-				if respContents != nil {
-					defer func() { _ = respContents.Body.Close() }()
+				gqlClient, err := getGQLClient(ctx)
+				if err != nil {
+					return mcp.NewToolResultError("failed to get GitHub GraphQL client"), nil
 				}
-				if err == nil && respContents.StatusCode == http.StatusOK && fileContent != nil && fileContent.SHA != nil {
-					fileSHA = *fileContent.SHA
+				fileSHA, err := getFileSHAFunc(ctx, gqlClient, owner, repo, path, rawOpts)
+				if err != nil {
+					return mcp.NewToolResultError(fmt.Sprintf("failed to get file SHA: %s", err)), nil
 				}
 
 				rawClient, err := getRawClient(ctx)
@@ -1383,3 +1386,34 @@ func resolveGitReference(ctx context.Context, githubClient *github.Client, owner
 	// Use provided ref, or it will be empty which defaults to the default branch
 	return &raw.ContentOpts{Ref: ref, SHA: sha}, nil
 }
+
+// getFileSHA retrieves the Blob SHA of a file.
+func getFileSHA(ctx context.Context, client *githubv4.Client, owner, repo, path string, opts *raw.ContentOpts) (string, error) {
+	var query struct {
+		Repository struct {
+			Object struct {
+				Blob struct {
+					OID string
+				} `graphql:"... on Blob"`
+			} `graphql:"object(expression: $expression)"`
+		} `graphql:"repository(owner: $owner, name: $repo)"`
+	}
+
+	// Prepare the expression based on the provided options
+	expression := githubv4.String(path)
+	if opts != nil && opts.SHA != "" {
+		expression = githubv4.String(fmt.Sprintf("%s:%s", opts.SHA, path))
+	} else if opts != nil && opts.Ref != "" {
+		expression = githubv4.String(fmt.Sprintf("%s:%s", opts.Ref, path))
+	}
+
+	variables := map[string]interface{}{
+		"owner":      githubv4.String(owner),
+		"repo":       githubv4.String(repo),
+		"expression": expression,
+	}
+	if err := client.Query(ctx, &query, variables); err != nil {
+		return "", fmt.Errorf("failed to query file SHA: %w", err)
+	}
+	return query.Repository.Object.Blob.OID, nil
+}

pkg/github/repositories_test.go

@@ -15,15 +15,24 @@ import (
 	"github.com/google/go-github/v72/github"
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/migueleliasweb/go-github-mock/src/mock"
+	"github.com/shurcooL/githubv4"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+// mockGetFileSHA is a test helper that mocks the getFileSHA function so that we don't need to mock the GraphQL client in Test_GetFileContents.
+func mockGetFileSHA(expectedSHA string) func(context.Context, *githubv4.Client, string, string, string, *raw.ContentOpts) (string, error) {
+	return func(_ context.Context, _ *githubv4.Client, _, _, _ string, _ *raw.ContentOpts) (string, error) {
+		return expectedSHA, nil
+	}
+}
+
 func Test_GetFileContents(t *testing.T) {
 	// Verify tool definition once
 	mockClient := github.NewClient(nil)
 	mockRawClient := raw.NewClient(mockClient, &url.URL{Scheme: "https", Host: "raw.githubusercontent.com", Path: "/"})
-	tool, _ := GetFileContents(stubGetClientFn(mockClient), stubGetRawClientFn(mockRawClient), translations.NullTranslationHelper)
+	mockGQLClient := githubv4.NewClient(nil)
+	tool, _ := GetFileContents(stubGetClientFn(mockClient), stubGetRawClientFn(mockRawClient), stubGetGQLClientFn(mockGQLClient), translations.NullTranslationHelper)
 	require.NoError(t, toolsnaps.Test(tool.Name, tool))
 
 	assert.Equal(t, "get_file_contents", tool.Name)
@@ -56,10 +65,10 @@ func Test_GetFileContents(t *testing.T) {
 			HTMLURL: github.Ptr("https://github.com/owner/repo/tree/main/src"),
 		},
 	}
-
 	tests := []struct {
 		name           string
 		mockedClient   *http.Client
+		mockFileSHA    string
 		requestArgs    map[string]interface{}
 		expectError    bool
 		expectedResult interface{}
@@ -98,6 +107,7 @@ func Test_GetFileContents(t *testing.T) {
 					}),
 				),
 			),
+			mockFileSHA: "abc123",
 			requestArgs: map[string]interface{}{
 				"owner": "owner",
 				"repo":  "repo",
@@ -143,6 +153,7 @@ func Test_GetFileContents(t *testing.T) {
 					}),
 				),
 			),
+			mockFileSHA: "def456",
 			requestArgs: map[string]interface{}{
 				"owner": "owner",
 				"repo":  "repo",
@@ -188,6 +199,7 @@ func Test_GetFileContents(t *testing.T) {
 					),
 				),
 			),
+			mockFileSHA: "", // Directory content doesn't need SHA
 			requestArgs: map[string]interface{}{
 				"owner": "owner",
 				"repo":  "repo",
@@ -221,6 +233,7 @@ func Test_GetFileContents(t *testing.T) {
 					}),
 				),
 			),
+			mockFileSHA: "", // Error case doesn't need SHA
 			requestArgs: map[string]interface{}{
 				"owner": "owner",
 				"repo":  "repo",
@@ -234,10 +247,20 @@ func Test_GetFileContents(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			// Mock the getFileSHA function if mockFileSHA is provided
+			if tc.mockFileSHA != "" {
+				originalGetFileSHA := getFileSHAFunc
+				getFileSHAFunc = mockGetFileSHA(tc.mockFileSHA)
+				defer func() {
+					getFileSHAFunc = originalGetFileSHA
+				}()
+			}
+
 			// Setup client with mock
 			client := github.NewClient(tc.mockedClient)
 			mockRawClient := raw.NewClient(client, &url.URL{Scheme: "https", Host: "raw.example.com", Path: "/"})
-			_, handler := GetFileContents(stubGetClientFn(client), stubGetRawClientFn(mockRawClient), translations.NullTranslationHelper)
+			mockGQLClient := githubv4.NewClient(tc.mockedClient)
+			_, handler := GetFileContents(stubGetClientFn(client), stubGetRawClientFn(mockRawClient), stubGetGQLClientFn(mockGQLClient), translations.NullTranslationHelper)
 
 			// Create call request
 			request := createMCPRequest(tc.requestArgs)

pkg/github/tools.go

@@ -24,7 +24,7 @@ func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetG
 	repos := toolsets.NewToolset("repos", "GitHub Repository related tools").
 		AddReadTools(
 			toolsets.NewServerTool(SearchRepositories(getClient, t)),
-			toolsets.NewServerTool(GetFileContents(getClient, getRawClient, t)),
+			toolsets.NewServerTool(GetFileContents(getClient, getRawClient, getGQLClient, t)),
 			toolsets.NewServerTool(ListCommits(getClient, t)),
 			toolsets.NewServerTool(SearchCode(getClient, t)),
 			toolsets.NewServerTool(GetCommit(getClient, t)),