diff --git a/.github/actions/amplify-preview/action.yaml b/.github/actions/amplify-preview/action.yaml
deleted file mode 100644
index e34e627..0000000
--- a/.github/actions/amplify-preview/action.yaml
+++ /dev/null
@@ -1,138 +0,0 @@
-name: Amplify Preview
-description: Prepare Amplify Preview URL and post it in PR comments
-inputs:
-  app_ids:
-    description: "Comma separated list of Amplify App IDs"
-    required: true
-  create_branches:
-    description: 'Create preview branches using this actions instead of "auto-build" feature on AWS Amplify'
-    required: false
-    default: "false"
-runs:
-  using: composite
-  steps:
-    - name: Extract branch name
-      shell: bash
-      run: echo "branch=${GITHUB_HEAD_REF:-${GITHUB_REF#refs/heads/}}" >> $GITHUB_OUTPUT
-      id: extract_branch
-
-    - name: Get or create Amplify branch
-      id: get_amplify_branch
-      env:
-        AMPLIFY_APP_IDS: ${{ inputs.app_ids }}
-        BRANCH_NAME: ${{ steps.extract_branch.outputs.branch }}
-        CREATE_BRANCHES: ${{ inputs.create_branches }}
-      shell: bash
-      run: |
-        set +e
-        github_outputs() {
-            echo "APP_ID=${1}"
-            echo "PREVIEW_URL=https://$(jq -r '.displayName' <<< "$2").$app_id.amplifyapp.com"
-            echo "CREATE_TIME=$(jq -r '.createTime' <<< "$2")"
-            echo "UPDATE_TIME=$(jq -r '.updateTime' <<< "$2")"
-            echo "JOB_ID=$(jq -r '.activeJobId' <<< "$2")"
-        }
-        echo "ℹ️ Checking if branch connected to any of the apps" >&2
-        IFS=, app_id_array=("$AMPLIFY_APP_IDS")
-        for app_id in "${app_id_array[@]}"; do
-          if branch_info=$(aws amplify get-branch --app-id "${app_id}" --branch-name "${BRANCH_NAME}" --query 'branch'); then
-            found_in_app_id=$app_id
-            github_outputs "$app_id" "$branch_info" >> "$GITHUB_OUTPUT"
-            echo "ℹ️ ${BRANCH_NAME} is already connected to app ${app_id}" >&2
-            break
-          fi
-        done
-        if [ -z "${found_in_app_id}" ] && [ "${CREATE_BRANCHES}" = "true" ]; then
-          echo "ℹ️ Trying to connect branch to one of the apps" >&2
-          for app_id in "${app_id_array[@]}"; do
-            if branch_info=$(aws amplify create-branch --app-id "${app_id}" --branch-name "${BRANCH_NAME}" --stage PULL_REQUEST --enable-auto-build --query 'branch'); then
-              found_in_app_id=$app_id
-              github_outputs "$app_id" "$branch_info" >> "$GITHUB_OUTPUT"
-              echo "✅ ${BRANCH_NAME} was successfully connected to app ${app_id}" >&2
-              break
-            fi
-          done
-        fi
-        if [ -z "${found_in_app_id}" ]; then
-          echo "❌ Branch wasn't connected to any of the apps" >&2
-          exit 1
-        fi
-
-    - name: Get or create Amplify deployment
-      id: get_amplify_job
-      env:
-        APP_ID: ${{ steps.get_amplify_branch.outputs.APP_ID }}
-        JOB_ID: ${{ steps.get_amplify_branch.outputs.JOB_ID }}
-        BRANCH_NAME: ${{ steps.extract_branch.outputs.branch }}
-        CREATE_BRANCHES: ${{ inputs.create_branches }}
-      shell: bash
-      continue-on-error: true
-      run: |
-        if [ "${CREATE_BRANCHES}" = "true" ] && { [ -z "${JOB_ID}" ] || [ "${JOB_ID}" = "null" ]; }; then
-          job_info=$(aws amplify start-job --app-id "${APP_ID}" --branch-name "${BRANCH_NAME}" --job-type RELEASE --job-reason "${GITHUB_WORKFLOW} (${GITHUB_RUN_ID}) - ${GITHUB_REF_NAME}" --query 'jobSummary')
-          if [ $? -eq 0 ]; then
-            JOB_ID=$(jq -r '.jobId' <<< "$job_info")
-            echo "🚀 Successfully triggered job ${JOB_ID} for ${BRANCH_NAME} on app ${APP_ID}" >&2
-          fi
-        fi
-
-        job_info=$(aws amplify get-job --app-id "${APP_ID}" --branch-name "${BRANCH_NAME}" --job-id "${JOB_ID}" --query 'job.summary')
-        {
-          echo "JOB_ID=$(jq -r '.jobId' <<< "$job_info")"
-          echo "JOB_STATUS=$(jq -r '.status' <<< "$job_info")"
-          echo "COMMIT_ID=$(jq -r '.commitId' <<< "$job_info")"
-        } >> "$GITHUB_OUTPUT"
-
-    - uses: actions/github-script@v7
-      if: ${{ github.event_name == 'pull_request' }}
-      env:
-        PREVIEW_URL: ${{ steps.get_amplify_branch.outputs.PREVIEW_URL }}
-        UPDATE_TIME: ${{ steps.get_amplify_branch.outputs.UPDATE_TIME }}
-        JOB_ID: ${{ steps.get_amplify_job.outputs.JOB_ID || steps.get_amplify_branch.outputs.JOB_ID }}
-        JOB_STATUS: ${{ steps.get_amplify_job.outputs.JOB_STATUS }}
-        COMMIT_ID: ${{ steps.get_amplify_job.outputs.COMMIT_ID }}
-      with:
-        script: |
-          const previewUrl = process.env.PREVIEW_URL;
-          const jobId = process.env.JOB_ID;
-          const jobStatus = process.env.JOB_STATUS || "unknown";
-          const updatedAt = process.env.UPDATE_TIME;
-          const commitId = process.env.COMMIT_ID;
-
-          const commentBody = `![🤖](https://a0.awsstatic.com/libra-css/images/site/fav/favicon.ico) Amplify preview here: ${previewUrl}
-
-          <details><summary>Preview details</summary>
-
-          - **LAST_UPDATED_AT**: ${updatedAt}
-          - **JOB_ID**: ${jobId}
-          - **JOB_STATUS**: ${jobStatus}
-          - **COMMIT_ID**: ${commitId}
-
-          </details>
-          `;
-
-          const prProps = {
-            issue_number: context.issue.number,
-            owner: context.repo.owner,
-            repo: context.repo.repo,
-          };
-
-          const comments = (await github.rest.issues.listComments(prProps))?.data;
-
-          const existingComment = comments?.find((comment) =>
-            comment.user.login === "github-actions[bot]" && comment.body.includes("Amplify preview here"));
-
-          if (existingComment == null) {
-            console.log("Posting new comment ${existingComment.id}")
-            github.rest.issues.createComment({
-              ...prProps,
-              body: commentBody,
-            })
-          } else {
-            console.log("Found existing comment ${existingComment.id}")
-            github.rest.issues.updateComment({
-              ...prProps,
-              body: commentBody,
-              comment_id: existingComment.id,
-            })
-          }
\ No newline at end of file
diff --git a/tools/amplify-preview/action.yaml b/tools/amplify-preview/action.yaml
index 260f59e..e53fc56 100644
--- a/tools/amplify-preview/action.yaml
+++ b/tools/amplify-preview/action.yaml
@@ -11,6 +11,9 @@ inputs:
   github_token:
     required: true
     description: "Github token with permissions to read/write comments in pull request"
+  wait:
+    default: "false"
+    description: "If Amplify deployment is pending/running state wait for it's completion"
 runs:
   using: composite
   steps:
@@ -30,6 +33,7 @@ runs:
         GIT_BRANCH_NAME: ${{ steps.extract_branch.outputs.branch }}
         CREATE_BRANCHES: ${{ inputs.create_branches }}
         GITHUB_TOKEN: ${{ inputs.github_token }}
+        WAIT: ${{ inputs.wait }} 
       shell: bash
       run: |
         pushd ./tools/amplify-preview/; go run ./; popd
diff --git a/tools/amplify-preview/amplify.go b/tools/amplify-preview/amplify.go
index 054cb16..20e2664 100644
--- a/tools/amplify-preview/amplify.go
+++ b/tools/amplify-preview/amplify.go
@@ -15,8 +15,13 @@ import (
 )
 
 var (
-	errBranchNotFound = errors.New("Branch not found")
-	errNoJobForBranch = errors.New("Current branch has no jobs")
+	errBranchNotFound           = errors.New("Branch not found")
+	errNoJobForBranch           = errors.New("Current branch has no jobs")
+	amplifyJobCompletedStatuses = map[types.JobStatus]struct{}{
+		types.JobStatusFailed:    {},
+		types.JobStatusCancelled: {},
+		types.JobStatusSucceed:   {},
+	}
 )
 
 const (
@@ -186,6 +191,11 @@ func appIDFromBranchARN(branchArn string) (string, error) {
 	return "", fmt.Errorf("Invalid branch ARN")
 }
 
+func isAmplifyJobCompleted(job *types.JobSummary) bool {
+	_, ok := amplifyJobCompletedStatuses[job.Status]
+	return ok
+}
+
 func (err aggregatedError) Error() error {
 	if len(err.perAppErr) == 0 {
 		return nil
diff --git a/tools/amplify-preview/main.go b/tools/amplify-preview/main.go
index f94c50d..9bbcee0 100644
--- a/tools/amplify-preview/main.go
+++ b/tools/amplify-preview/main.go
@@ -22,6 +22,7 @@ import (
 	"log/slog"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -38,6 +39,13 @@ var (
 	gitBranchName  = kingpin.Flag("git-branch-name", "Git branch name").Envar("GIT_BRANCH_NAME").Required().String()
 	createBranches = kingpin.Flag("crate-branches",
 		"Defines whether Amplify branches should be created if missing, or just lookup existing ones").Envar("CREATE_BRANCHES").Default("false").Bool()
+	wait = kingpin.Flag("wait",
+		"Wait for pending/running job to complete").Envar("wait").Default("false").Bool()
+)
+
+const (
+	jobWaitSleepTime    = 30 * time.Second
+	jobWaitTimeAttempts = 40
 )
 
 func main() {
@@ -96,6 +104,17 @@ func main() {
 		}
 	}
 
+	for i := 0; *wait && !isAmplifyJobCompleted(job) && i < jobWaitTimeAttempts; i++ {
+		job, err := amp.GetJob(ctx, branch, nil)
+		if err != nil {
+			logger.Error("failed to get amplify job", logKeyBranchName, *gitBranchName, "error", err)
+			os.Exit(1)
+		}
+
+		logger.Info("Job is not in a completed state yet. Sleeping...", logKeyBranchName, *gitBranchName, "job_status", job.Status, "job_id", job.JobId)
+		time.Sleep(jobWaitSleepTime)
+	}
+
 	if err := postPreviewURL(ctx, amplifyJobToMarkdown(job, branch)); err != nil {
 		logger.Error("failed to post preview URL", "error", err)
 		os.Exit(1)
@@ -104,7 +123,7 @@ func main() {
 	logger.Info("Successfully posted PR comment")
 
 	if job.Status == types.JobStatusFailed {
-		logger.Error("amplify job is in failed state", "job_status", job.Status, "job_id", job.JobId)
+		logger.Error("amplify job is in failed state", logKeyBranchName, *gitBranchName, "job_status", job.Status, "job_id", job.JobId)
 		os.Exit(1)
 	}
 }