diff --git a/.github/workflows/build-image.yml b/.github/workflows/build-image.yml
index 3bfc7b2..2b94699 100644
--- a/.github/workflows/build-image.yml
+++ b/.github/workflows/build-image.yml
@@ -21,10 +21,11 @@ jobs:
           echo "date=$(date --rfc-3339=seconds | sed 's/ /T/')" >> $GITHUB_OUTPUT
 
   version:
-    name: Get latest package version
+    name: Determine package version and image tag
     runs-on: ubuntu-latest
     outputs:
       package_version: ${{ steps.version.outputs.package_version }}
+      image_tag: ${{ steps.version.outputs.image_tag }}
     steps:
       - name: Add Microsoft's APT repo for VS Code
         run: |
@@ -36,12 +37,16 @@ jobs:
             | sudo tee -a /etc/apt/sources.list.d/microsoft.list
           sudo apt-get update -y
 
-      - name: Get package version
+      - name: Get package version and set image tag 
         id: version
         run: |
-          package_version=$(apt list code 2>/dev/null | grep -v "^Listing" | awk '{print $2}' | cut -d'-' -f1)
-          echo "package_version=${package_version}" >> $GITHUB_OUTPUT
-          echo "The package version is $package_version"
+          package_version=$(apt-cache policy code | awk -F': ' '/Candidate:/ {print $2}')
+          [ -n "$package_version" ] || { echo "Failed to determine VS Code package version" >&2; exit 1; }
+          image_tag=$(echo "$package_version" | cut -d'-' -f1)
+          echo "package_version=$package_version" >> $GITHUB_OUTPUT
+          echo "image_tag=$image_tag" >> $GITHUB_OUTPUT
+          echo "Package Version: $package_version"
+          echo "Image Tag: $image_tag"
 
   build:
     name: Build and push image
@@ -79,11 +84,12 @@ jobs:
           push: true
           build-args: |
             IMAGE_BUILD_DATE=${{ needs.date.outputs.date }}
+            INSTALL_VERSION=${{ needs.version.outputs.package_version }}
           tags: |
             ghcr.io/${{ github.repository_owner }}/code-tunnel:latest
-            ghcr.io/${{ github.repository_owner }}/code-tunnel:${{ needs.version.outputs.package_version }}
+            ghcr.io/${{ github.repository_owner }}/code-tunnel:${{ needs.version.outputs.image_tag }}
             ${{ vars.DOCKERHUB_USERNAME }}/code-tunnel:latest
-            ${{ vars.DOCKERHUB_USERNAME }}/code-tunnel:${{ needs.version.outputs.package_version }}
+            ${{ vars.DOCKERHUB_USERNAME }}/code-tunnel:${{ needs.version.outputs.image_tag }}
           labels: |
             org.opencontainers.image.ref.name=${{ github.ref }}
             org.opencontainers.image.revision=${{ github.sha }}
diff --git a/Dockerfile b/Dockerfile
index 625bfb6..c398821 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,6 +2,7 @@ FROM ghcr.io/linuxserver/baseimage-ubuntu:noble
 
 # set labels
 ARG IMAGE_BUILD_DATE
+ARG INSTALL_VERSION
 LABEL release_channel="stable"
 LABEL org.opencontainers.image.authors="tibynx"
 LABEL org.opencontainers.image.created="${IMAGE_BUILD_DATE}"
@@ -48,7 +49,7 @@ RUN \
     https://packages.microsoft.com/repos/code stable main" > /etc/apt/sources.list.d/vscode.list && \
   apt-get update -y && \
   apt-get install --no-install-recommends -y \
-    code && \
+    code=${INSTALL_VERSION} && \
   echo "**** clean up ****" && \
   apt-get clean && \
   rm -rf \
diff --git a/Dockerfile.aarch64 b/Dockerfile.aarch64
index 8c589b1..b0e0c29 100644
--- a/Dockerfile.aarch64
+++ b/Dockerfile.aarch64
@@ -2,6 +2,7 @@ FROM ghcr.io/linuxserver/baseimage-ubuntu:arm64v8-noble
 
 # set labels
 ARG IMAGE_BUILD_DATE
+ARG INSTALL_VERSION
 LABEL release_channel="stable"
 LABEL org.opencontainers.image.authors="tibynx"
 LABEL org.opencontainers.image.created="${IMAGE_BUILD_DATE}"
@@ -48,7 +49,7 @@ RUN \
     https://packages.microsoft.com/repos/code stable main" > /etc/apt/sources.list.d/vscode.list && \
   apt-get update -y && \
   apt-get install --no-install-recommends -y \
-    code && \
+    code=${INSTALL_VERSION} && \
   echo "**** clean up ****" && \
   apt-get clean && \
   rm -rf \