Use REST endpoints for PR draft state transitions

pkg/github/granular_tools_test.go

@@ -1499,89 +1499,67 @@ func TestGranularCreatePullRequestReview(t *testing.T) {
 
 func TestGranularUpdatePullRequestDraftState(t *testing.T) {
 	tests := []struct {
-		name  string
-		draft bool
+		name                string
+		initialDraftState   bool
+		requestedDraftState bool
+		expectedActionCalls int
+		getStatusCode       int
+		actionStatusCode    int
+		expectError         bool
+		expectedErrContains string
 	}{
-		{name: "convert to draft", draft: true},
-		{name: "mark ready for review", draft: false},
+		{name: "convert to draft", initialDraftState: false, requestedDraftState: true, expectedActionCalls: 1, getStatusCode: http.StatusOK, actionStatusCode: http.StatusOK},
+		{name: "mark ready for review", initialDraftState: true, requestedDraftState: false, expectedActionCalls: 1, getStatusCode: http.StatusOK, actionStatusCode: http.StatusOK},
+		{name: "no-op when already requested state", initialDraftState: true, requestedDraftState: true, expectedActionCalls: 0, getStatusCode: http.StatusOK, actionStatusCode: http.StatusOK},
+		{name: "returns error when action endpoint is forbidden", initialDraftState: false, requestedDraftState: true, expectedActionCalls: 1, getStatusCode: http.StatusOK, actionStatusCode: http.StatusForbidden, expectError: true, expectedErrContains: "failed to convert pull request to draft"},
+		{name: "returns error when pull request lookup is not found", initialDraftState: false, requestedDraftState: true, expectedActionCalls: 0, getStatusCode: http.StatusNotFound, actionStatusCode: http.StatusOK, expectError: true, expectedErrContains: "failed to get pull request"},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			var matchers []githubv4mock.Matcher
-
-			matchers = append(matchers, githubv4mock.NewQueryMatcher(
-				struct {
-					Repository struct {
-						PullRequest struct {
-							ID githubv4.ID
-						} `graphql:"pullRequest(number: $number)"`
-					} `graphql:"repository(owner: $owner, name: $name)"`
-				}{},
-				map[string]any{
-					"owner":  githubv4.String("owner"),
-					"name":   githubv4.String("repo"),
-					"number": githubv4.Int(1),
+			actionCalls := 0
+			client := mustNewGHClient(t, MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+				GetReposPullsByOwnerByRepoByPullNumber: func(w http.ResponseWriter, _ *http.Request) {
+					if tc.getStatusCode != http.StatusOK {
+						writeJSONResponse(t, w, tc.getStatusCode, map[string]any{"message": "not found"})
+						return
+					}
+					writeJSONResponse(t, w, http.StatusOK, &gogithub.PullRequest{
+						Number: gogithub.Ptr(1),
+						Draft:  gogithub.Ptr(tc.initialDraftState),
+					})
 				},
-				githubv4mock.DataResponse(map[string]any{
-					"repository": map[string]any{
-						"pullRequest": map[string]any{"id": "PR_123"},
-					},
-				}),
-			))
-
-			if tc.draft {
-				matchers = append(matchers, githubv4mock.NewMutationMatcher(
-					struct {
-						ConvertPullRequestToDraft struct {
-							PullRequest struct {
-								ID      githubv4.ID
-								IsDraft githubv4.Boolean
-							}
-						} `graphql:"convertPullRequestToDraft(input: $input)"`
-					}{},
-					githubv4.ConvertPullRequestToDraftInput{PullRequestID: githubv4.ID("PR_123")},
-					nil,
-					githubv4mock.DataResponse(map[string]any{
-						"convertPullRequestToDraft": map[string]any{
-							"pullRequest": map[string]any{"id": "PR_123", "isDraft": true},
-						},
-					}),
-				))
-			} else {
-				matchers = append(matchers, githubv4mock.NewMutationMatcher(
-					struct {
-						MarkPullRequestReadyForReview struct {
-							PullRequest struct {
-								ID      githubv4.ID
-								IsDraft githubv4.Boolean
-							}
-						} `graphql:"markPullRequestReadyForReview(input: $input)"`
-					}{},
-					githubv4.MarkPullRequestReadyForReviewInput{PullRequestID: githubv4.ID("PR_123")},
-					nil,
-					githubv4mock.DataResponse(map[string]any{
-						"markPullRequestReadyForReview": map[string]any{
-							"pullRequest": map[string]any{"id": "PR_123", "isDraft": false},
-						},
-					}),
-				))
-			}
+				PostReposPullsByOwnerByRepoByPullNumberConvertToDraft: func(w http.ResponseWriter, _ *http.Request) {
+					actionCalls++
+					writeJSONResponse(t, w, tc.actionStatusCode, map[string]any{"message": "forbidden"})
+				},
+				PostReposPullsByOwnerByRepoByPullNumberReadyForReview: func(w http.ResponseWriter, _ *http.Request) {
+					actionCalls++
+					writeJSONResponse(t, w, tc.actionStatusCode, map[string]any{"message": "forbidden"})
+				},
+			}))
 
-			gqlClient := githubv4.NewClient(githubv4mock.NewMockedHTTPClient(matchers...))
-			deps := BaseDeps{GQLClient: gqlClient}
+			deps := BaseDeps{Client: client}
 			serverTool := GranularUpdatePullRequestDraftState(translations.NullTranslationHelper)
 			handler := serverTool.Handler(deps)
 
 			request := createMCPRequest(map[string]any{
 				"owner":      "owner",
 				"repo":       "repo",
 				"pullNumber": float64(1),
-				"draft":      tc.draft,
+				"draft":      tc.requestedDraftState,
 			})
 			result, err := handler(ContextWithDeps(context.Background(), deps), &request)
 			require.NoError(t, err)
+			if tc.expectError {
+				assert.True(t, result.IsError)
+				errorContent := getErrorResult(t, result)
+				assert.Contains(t, errorContent.Text, tc.expectedErrContains)
+				assert.Equal(t, tc.expectedActionCalls, actionCalls)
+				return
+			}
 			assert.False(t, result.IsError)
+			assert.Equal(t, tc.expectedActionCalls, actionCalls)
 		})
 	}
 }

pkg/github/helper_test.go

@@ -81,6 +81,8 @@ const (
 	PatchReposPullsByOwnerByRepoByPullNumber                  = "PATCH /repos/{owner}/{repo}/pulls/{pull_number}"
 	PutReposPullsMergeByOwnerByRepoByPullNumber               = "PUT /repos/{owner}/{repo}/pulls/{pull_number}/merge"
 	PutReposPullsUpdateBranchByOwnerByRepoByPullNumber        = "PUT /repos/{owner}/{repo}/pulls/{pull_number}/update-branch"
+	PostReposPullsByOwnerByRepoByPullNumberConvertToDraft     = "POST /repos/{owner}/{repo}/pulls/{pull_number}/convert-to-draft"
+	PostReposPullsByOwnerByRepoByPullNumberReadyForReview     = "POST /repos/{owner}/{repo}/pulls/{pull_number}/ready_for_review"
 	PostReposPullsRequestedReviewersByOwnerByRepoByPullNumber = "POST /repos/{owner}/{repo}/pulls/{pull_number}/requested_reviewers"
 	PostReposPullsCommentsByOwnerByRepoByPullNumber           = "POST /repos/{owner}/{repo}/pulls/{pull_number}/comments"
 	PostReposPullsCommentsReactionsByOwnerByRepoByCommentID   = "POST /repos/{owner}/{repo}/pulls/comments/{comment_id}/reactions"

pkg/github/pullrequests.go

@@ -23,6 +23,60 @@ import (
 	"github.com/github/github-mcp-server/pkg/utils"
 )
 
+func updatePullRequestDraftState(ctx context.Context, deps ToolDependencies, owner, repo string, pullNumber int, draft bool) *mcp.CallToolResult {
+	client, err := deps.GetClient(ctx)
+	if err != nil {
+		return utils.NewToolResultErrorFromErr("failed to get GitHub client", err)
+	}
+
+	pullRequest, getResp, err := client.PullRequests.Get(ctx, owner, repo, pullNumber)
+	if err != nil {
+		return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to get pull request", getResp, err)
+	}
+	defer func() { _ = getResp.Body.Close() }()
+
+	if getResp.StatusCode != http.StatusOK {
+		bodyBytes, err := io.ReadAll(getResp.Body)
+		if err != nil {
+			return utils.NewToolResultErrorFromErr("failed to read response body", err)
+		}
+		return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to get pull request", getResp, bodyBytes)
+	}
+
+	if pullRequest.GetDraft() == draft {
+		return nil
+	}
+
+	actionPath := "ready_for_review"
+	actionName := "mark pull request ready for review"
+	if draft {
+		actionPath = "convert-to-draft"
+		actionName = "convert pull request to draft"
+	}
+
+	apiURL := fmt.Sprintf("repos/%s/%s/pulls/%d/%s", owner, repo, pullNumber, actionPath)
+	req, err := client.NewRequest(ctx, http.MethodPost, apiURL, nil)
+	if err != nil {
+		return utils.NewToolResultErrorFromErr("failed to create request", err)
+	}
+
+	actionResp, err := client.Do(req, nil)
+	if err != nil {
+		return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to "+actionName, actionResp, err)
+	}
+	defer func() { _ = actionResp.Body.Close() }()
+
+	if actionResp.StatusCode != http.StatusOK {
+		bodyBytes, err := io.ReadAll(actionResp.Body)
+		if err != nil {
+			return utils.NewToolResultErrorFromErr("failed to read response body", err)
+		}
+		return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to "+actionName, actionResp, bodyBytes)
+	}
+
+	return nil
+}
+
 // PullRequestRead creates a tool to get details of a specific pull request.
 func PullRequestRead(t translations.TranslationHelperFunc) inventory.ServerTool {
 	schema := &jsonschema.Schema{
@@ -1012,69 +1066,10 @@ func UpdatePullRequest(t translations.TranslationHelperFunc) inventory.ServerToo
 				}
 			}
 
-			// Handle draft status changes using GraphQL
+			// Handle draft status changes using REST
 			if draftProvided {
-				gqlClient, err := deps.GetGQLClient(ctx)
-				if err != nil {
-					return utils.NewToolResultErrorFromErr("failed to get GitHub GraphQL client", err), nil, nil
-				}
-
-				var prQuery struct {
-					Repository struct {
-						PullRequest struct {
-							ID      githubv4.ID
-							IsDraft githubv4.Boolean
-						} `graphql:"pullRequest(number: $prNum)"`
-					} `graphql:"repository(owner: $owner, name: $repo)"`
-				}
-
-				err = gqlClient.Query(ctx, &prQuery, map[string]any{
-					"owner": githubv4.String(owner),
-					"repo":  githubv4.String(repo),
-					"prNum": githubv4.Int(pullNumber), // #nosec G115 - pull request numbers are always small positive integers
-				})
-				if err != nil {
-					return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "Failed to find pull request", err), nil, nil
-				}
-
-				currentIsDraft := bool(prQuery.Repository.PullRequest.IsDraft)
-
-				if currentIsDraft != draftValue {
-					if draftValue {
-						// Convert to draft
-						var mutation struct {
-							ConvertPullRequestToDraft struct {
-								PullRequest struct {
-									ID      githubv4.ID
-									IsDraft githubv4.Boolean
-								}
-							} `graphql:"convertPullRequestToDraft(input: $input)"`
-						}
-
-						err = gqlClient.Mutate(ctx, &mutation, githubv4.ConvertPullRequestToDraftInput{
-							PullRequestID: prQuery.Repository.PullRequest.ID,
-						}, nil)
-						if err != nil {
-							return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "Failed to convert pull request to draft", err), nil, nil
-						}
-					} else {
-						// Mark as ready for review
-						var mutation struct {
-							MarkPullRequestReadyForReview struct {
-								PullRequest struct {
-									ID      githubv4.ID
-									IsDraft githubv4.Boolean
-								}
-							} `graphql:"markPullRequestReadyForReview(input: $input)"`
-						}
-
-						err = gqlClient.Mutate(ctx, &mutation, githubv4.MarkPullRequestReadyForReviewInput{
-							PullRequestID: prQuery.Repository.PullRequest.ID,
-						}, nil)
-						if err != nil {
-							return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "Failed to mark pull request ready for review", err), nil, nil
-						}
-					}
+				if result := updatePullRequestDraftState(ctx, deps, owner, repo, pullNumber, draftValue); result != nil {
+					return result, nil, nil
 				}
 			}
 

pkg/github/pullrequests_granular.go

@@ -15,7 +15,6 @@ import (
 	gogithub "github.com/google/go-github/v87/github"
 	"github.com/google/jsonschema-go/jsonschema"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
-	"github.com/shurcooL/githubv4"
 )
 
 // prUpdateTool is a helper to create single-field pull request update tools via REST.
@@ -218,57 +217,14 @@ func GranularUpdatePullRequestDraftState(t translations.TranslationHelperFunc) i
 				return utils.NewToolResultError(err.Error()), nil, nil
 			}
 
-			gqlClient, err := deps.GetGQLClient(ctx)
-			if err != nil {
-				return utils.NewToolResultErrorFromErr("failed to get GitHub GraphQL client", err), nil, nil
-			}
-
-			// Get PR node ID
-			var prQuery struct {
-				Repository struct {
-					PullRequest struct {
-						ID githubv4.ID
-					} `graphql:"pullRequest(number: $number)"`
-				} `graphql:"repository(owner: $owner, name: $name)"`
-			}
-			if err := gqlClient.Query(ctx, &prQuery, map[string]any{
-				"owner":  githubv4.String(owner),
-				"name":   githubv4.String(repo),
-				"number": githubv4.Int(pullNumber), // #nosec G115 - PR numbers are always small positive integers
-			}); err != nil {
-				return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "failed to get pull request", err), nil, nil
+			if result := updatePullRequestDraftState(ctx, deps, owner, repo, pullNumber, draft); result != nil {
+				return result, nil, nil
 			}
 
 			if draft {
-				var mutation struct {
-					ConvertPullRequestToDraft struct {
-						PullRequest struct {
-							ID      githubv4.ID
-							IsDraft githubv4.Boolean
-						}
-					} `graphql:"convertPullRequestToDraft(input: $input)"`
-				}
-				if err := gqlClient.Mutate(ctx, &mutation, githubv4.ConvertPullRequestToDraftInput{
-					PullRequestID: prQuery.Repository.PullRequest.ID,
-				}, nil); err != nil {
-					return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "failed to convert to draft", err), nil, nil
-				}
 				return utils.NewToolResultText("pull request converted to draft"), nil, nil
 			}
 
-			var mutation struct {
-				MarkPullRequestReadyForReview struct {
-					PullRequest struct {
-						ID      githubv4.ID
-						IsDraft githubv4.Boolean
-					}
-				} `graphql:"markPullRequestReadyForReview(input: $input)"`
-			}
-			if err := gqlClient.Mutate(ctx, &mutation, githubv4.MarkPullRequestReadyForReviewInput{
-				PullRequestID: prQuery.Repository.PullRequest.ID,
-			}, nil); err != nil {
-				return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "failed to mark ready for review", err), nil, nil
-			}
 			return utils.NewToolResultText("pull request marked as ready for review"), nil, nil
 		},
 	)