diff --git a/scripts/ci/codegen/__tests__/codegen.test.ts b/scripts/ci/codegen/__tests__/codegen.test.ts
index ef09391967..b4aa42c3cc 100644
--- a/scripts/ci/codegen/__tests__/codegen.test.ts
+++ b/scripts/ci/codegen/__tests__/codegen.test.ts
@@ -33,7 +33,7 @@ describe('codegen', () => {
         // @ts-expect-error a parameter is required
         upsertGenerationComment()
       ).rejects.toThrow(
-        '`upsertGenerationComment` requires a `trigger` parameter (`codegen` | `notification`).'
+        "'upsertGenerationComment' requires a 'trigger' parameter (notification | codegen | noGen | cleanup)."
       );
     });
 
diff --git a/scripts/ci/codegen/pushGeneratedCode.ts b/scripts/ci/codegen/pushGeneratedCode.ts
index 18c5894ba7..19f9be4e6a 100644
--- a/scripts/ci/codegen/pushGeneratedCode.ts
+++ b/scripts/ci/codegen/pushGeneratedCode.ts
@@ -80,7 +80,6 @@ Co-authored-by: %an <%ae>
   await run(`git push origin ${branchToPush}`);
 
   if (PR_NUMBER) {
-    await run(`git checkout ${baseBranch}`);
     await run(`yarn workspace scripts upsertGenerationComment codegen`);
   }
 }
diff --git a/scripts/ci/codegen/upsertGenerationComment.ts b/scripts/ci/codegen/upsertGenerationComment.ts
index 28c115d03e..b5cc7c0009 100644
--- a/scripts/ci/codegen/upsertGenerationComment.ts
+++ b/scripts/ci/codegen/upsertGenerationComment.ts
@@ -14,6 +14,7 @@ const allowedTriggers = ['notification', 'codegen', 'noGen', 'cleanup'];
 type Trigger = keyof typeof commentText;
 
 export async function getCommentBody(trigger: Trigger): Promise<string> {
+  // All of the case where we are not pushing generated code.
   if (
     trigger === 'notification' ||
     trigger === 'noGen' ||
@@ -24,10 +25,11 @@ export async function getCommentBody(trigger: Trigger): Promise<string> {
 ${commentText[trigger].body}`;
   }
 
-  const baseBranch = await run('git branch --show-current');
+  // We are on a codegen step on a pull request here
+  const generatedBranch = await run('git branch --show-current');
+  const baseBranch = generatedBranch.replace('generated/', '');
   const baseCommit = await run(`git show ${baseBranch} -s --format=%H`);
 
-  const generatedBranch = `generated/${baseBranch}`;
   const generatedCommit = await run(
     `git show ${generatedBranch} -s --format=%H`
   );
@@ -48,7 +50,9 @@ ${commentText.codegen.body(
 export async function upsertGenerationComment(trigger: Trigger): Promise<void> {
   if (!trigger || allowedTriggers.includes(trigger) === false) {
     throw new Error(
-      '`upsertGenerationComment` requires a `trigger` parameter (`codegen` | `notification`).'
+      `'upsertGenerationComment' requires a 'trigger' parameter (${allowedTriggers.join(
+        ' | '
+      )}).`
     );
   }