diff --git a/src/shell/scripts/omp.ps1 b/src/shell/scripts/omp.ps1
index 90118276..7e1f82b9 100644
--- a/src/shell/scripts/omp.ps1
+++ b/src/shell/scripts/omp.ps1
@@ -387,8 +387,6 @@ Example:
             }
             catch {}
         }
-
-        Set-PoshContext $script:ErrorCode
     }
 
     function Update-PoshErrorCode {
@@ -446,6 +444,8 @@ Example:
             Update-PoshErrorCode
         }
 
+        Set-PoshContext $script:ErrorCode
+
         $cleanPSWD = Get-CleanPSWD
         $stackCount = global:Get-PoshStackCount
         $terminalWidth = Get-TerminalWidth