3 files changed, 15 insertions, 41 deletions
diff --git a/src/github-api-helper.ts b/src/github-api-helper.ts
index 2124a86..86632f5 100644
--- a/src/github-api-helper.ts
+++ b/src/github-api-helper.ts
@@ -1,12 +1,12 @@
 import * as assert from 'assert'
 import * as core from '@actions/core'
 import * as fs from 'fs'
+import * as github from '@actions/github'
 import * as io from '@actions/io'
 import * as path from 'path'
 import * as retryHelper from './retry-helper'
 import * as toolCache from '@actions/tool-cache'
 import {default as uuid} from 'uuid/v4'
-import {getOctokit, Octokit} from './octokit-provider'
 
 const IS_WINDOWS = process.platform === 'win32'
 
@@ -84,11 +84,11 @@ export async function getDefaultBranch(
 ): Promise<string> {
   return await retryHelper.execute(async () => {
     core.info('Retrieving the default branch name')
-    const octokit = getOctokit(authToken, {baseUrl: baseUrl})
+    const octokit = github.getOctokit(authToken, {baseUrl: baseUrl})
     let result: string
     try {
       // Get the default branch from the repo info
-      const response = await octokit.repos.get({owner, repo})
+      const response = await octokit.rest.repos.get({owner, repo})
       result = response.data.default_branch
       assert.ok(result, 'default_branch cannot be empty')
     } catch (err) {
@@ -125,19 +125,14 @@ async function downloadArchive(
   commit: string,
   baseUrl?: string
 ): Promise<Buffer> {
-  const octokit = getOctokit(authToken, {baseUrl: baseUrl})
-  const params: Octokit.ReposGetArchiveLinkParams = {
+  const octokit = github.getOctokit(authToken, {baseUrl: baseUrl})
+  const download = IS_WINDOWS
+    ? octokit.rest.repos.downloadZipballArchive
+    : octokit.rest.repos.downloadTarballArchive
+  const response = await download({
     owner: owner,
     repo: repo,
-    archive_format: IS_WINDOWS ? 'zipball' : 'tarball',
     ref: commit || ref
-  }
-  const response = await octokit.repos.getArchiveLink(params)
-  if (response.status != 200) {
-    throw new Error(
-      `Unexpected response from GitHub API. Status: ${response.status}, Data: ${response.data}`
-    )
-  }
-
-  return Buffer.from(response.data) // response.data is ArrayBuffer
+  })
+  return Buffer.from(response.data as ArrayBuffer) // response.data is ArrayBuffer
 }
diff --git a/src/octokit-provider.ts b/src/octokit-provider.ts
deleted file mode 100644
index 4fb1285..0000000
--- a/src/octokit-provider.ts
+++ /dev/null
@@ -1,23 +0,0 @@
-import * as github from '@actions/github'
-import {Octokit} from '@octokit/rest'
-import {getServerApiUrl} from './url-helper'
-
-// Centralize all Octokit references by re-exporting
-export {Octokit} from '@octokit/rest'
-
-export type OctokitOptions = {
-  baseUrl?: string
-  userAgent?: string
-}
-
-export function getOctokit(authToken: string, opts: OctokitOptions) {
-  const options: Octokit.Options = {
-    baseUrl: getServerApiUrl(opts.baseUrl)
-  }
-
-  if (opts.userAgent) {
-    options.userAgent = opts.userAgent
-  }
-
-  return new github.GitHub(authToken, options)
-}
diff --git a/src/ref-helper.ts b/src/ref-helper.ts
index ce875fc..9b1ca95 100644
--- a/src/ref-helper.ts
+++ b/src/ref-helper.ts
@@ -1,7 +1,6 @@
 import {IGitCommandManager} from './git-command-manager'
 import * as core from '@actions/core'
 import * as github from '@actions/github'
-import {getOctokit} from './octokit-provider'
 import {isGhes} from './url-helper'
 
 export const tagsRefSpec = '+refs/tags/*:refs/tags/*'
@@ -245,7 +244,7 @@ export async function checkCommitInfo(
       core.debug(
         `Expected head sha ${expectedHeadSha}; actual head sha ${actualHeadSha}`
       )
-      const octokit = getOctokit(token, {
+      const octokit = github.getOctokit(token, {
         baseUrl: baseUrl,
         userAgent: `actions-checkout-tracepoint/1.0 (code=STALE_MERGE;owner=${repositoryOwner};repo=${repositoryName};pr=${fromPayload(
           'number'
@@ -253,7 +252,10 @@ export async function checkCommitInfo(
           process.env['GITHUB_RUN_ID']
         };expected_head_sha=${expectedHeadSha};actual_head_sha=${actualHeadSha})`
       })
-      await octokit.repos.get({owner: repositoryOwner, repo: repositoryName})
+      await octokit.rest.repos.get({
+        owner: repositoryOwner,
+        repo: repositoryName
+      })
     }
   } catch (err) {
     core.debug(