2 files changed, 9 insertions, 4 deletions
diff --git a/src/github-api-helper.ts b/src/github-api-helper.ts
index 86632f5..b6bfc13 100644
--- a/src/github-api-helper.ts
+++ b/src/github-api-helper.ts
@@ -7,6 +7,7 @@ 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 {getServerApiUrl} from './url-helper'
 
 const IS_WINDOWS = process.platform === 'win32'
 
@@ -84,7 +85,9 @@ export async function getDefaultBranch(
 ): Promise<string> {
   return await retryHelper.execute(async () => {
     core.info('Retrieving the default branch name')
-    const octokit = github.getOctokit(authToken, {baseUrl: baseUrl})
+    const octokit = github.getOctokit(authToken, {
+      baseUrl: getServerApiUrl(baseUrl)
+    })
     let result: string
     try {
       // Get the default branch from the repo info
@@ -125,7 +128,9 @@ async function downloadArchive(
   commit: string,
   baseUrl?: string
 ): Promise<Buffer> {
-  const octokit = github.getOctokit(authToken, {baseUrl: baseUrl})
+  const octokit = github.getOctokit(authToken, {
+    baseUrl: getServerApiUrl(baseUrl)
+  })
   const download = IS_WINDOWS
     ? octokit.rest.repos.downloadZipballArchive
     : octokit.rest.repos.downloadTarballArchive
diff --git a/src/ref-helper.ts b/src/ref-helper.ts
index 9b1ca95..1c25049 100644
--- a/src/ref-helper.ts
+++ b/src/ref-helper.ts
@@ -1,7 +1,7 @@
 import {IGitCommandManager} from './git-command-manager'
 import * as core from '@actions/core'
 import * as github from '@actions/github'
-import {isGhes} from './url-helper'
+import {getServerApiUrl, isGhes} from './url-helper'
 
 export const tagsRefSpec = '+refs/tags/*:refs/tags/*'
 
@@ -245,7 +245,7 @@ export async function checkCommitInfo(
         `Expected head sha ${expectedHeadSha}; actual head sha ${actualHeadSha}`
       )
       const octokit = github.getOctokit(token, {
-        baseUrl: baseUrl,
+        baseUrl: getServerApiUrl(baseUrl),
         userAgent: `actions-checkout-tracepoint/1.0 (code=STALE_MERGE;owner=${repositoryOwner};repo=${repositoryName};pr=${fromPayload(
           'number'
         )};run_id=${