From a9de502db13ebe4229224e81101c02e2ea126967 Mon Sep 17 00:00:00 2001
From: Shantanu <joshishantanu4@gmail.com>
Date: Thu, 25 Apr 2024 15:58:29 -1000
Subject: [PATCH] cmd/run: savvy run works for zsh

---
 cmd/setup/savvy.zsh | 25 +++++++++++++------------
 shell/zsh.go        |  9 ++++++++-
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/cmd/setup/savvy.zsh b/cmd/setup/savvy.zsh
index a22b4ed..b7457ef 100755
--- a/cmd/setup/savvy.zsh
+++ b/cmd/setup/savvy.zsh
@@ -2,6 +2,7 @@
 SAVVY_INPUT_FILE=/tmp/savvy-socket
 
 autoload -Uz add-zsh-hook
+autoload -Uz add-zle-hook-widget
 
 step_id=""
 
@@ -32,31 +33,31 @@ function __savvy_record_pre_exec__() {
 function __savvy_run_pre_exec__() {
   local cmd=$3
   if [[ "${SAVVY_CONTEXT}" == "run" ]] ; then
-    if [[ "${cmd}" -eq "${SAVVY_COMMANDS[SAVVY_NEXT_STEP]}" ]] ; then
-      export SAVVY_NEXT_STEP=$((SAVVY_NEXT_STEP+1))
+    if [[ "${cmd}" == "${SAVVY_COMMANDS[SAVVY_NEXT_STEP]}" ]] ; then
+      SAVVY_NEXT_STEP=$((SAVVY_NEXT_STEP+1))
     fi
   fi
 }
 
-# SAVVY_RUNBOOK_COMMANDS is a list of commands that savvy should run in the run context
-
-SAVVY_COMMANDS=("${(s:,:)SAVVY_RUNBOOK_COMMANDS}")
-num_commands=${#SAVVY_COMMANDS}
 function __savvy_runbook_runner__() {
-  next_step_idx=${SAVVY_NEXT_STEP:1}
-  BUFFER=${SAVVY_COMMANDS[next_step_idx]}  # Initial text to be edited by the user
-  zle end-of-line  # Accept the line for editing
+  if [[ "${SAVVY_CONTEXT}" == "run" ]] ; then
+    next_step_idx=${SAVVY_NEXT_STEP}
+    BUFFER="${SAVVY_COMMANDS[next_step_idx]}"
+    zle end-of-line  # Accept the line for editing
+  fi
 }
 
-
-
 # NOTE: If you change any function names, you must also change the corresponding check in shell/check_setup.go, shell/zsh.go
 #
 # TODO: use templates to avoid the need to manually change shell checks
 
+SAVVY_COMMANDS=()
+SAVVY_NEXT_STEP=1
 if [[ "${SAVVY_CONTEXT}" == "run" ]] ; then
   zle -N zle-line-init __savvy_runbook_runner__
-  add-zle-hook-widget zle-line-init
+  add-zle-hook-widget line-init __savvy_runbook_runner__
+  # SAVVY_RUNBOOK_COMMANDS is a list of commands that savvy should run in the run context
+  SAVVY_COMMANDS=("${(s:,:)SAVVY_RUNBOOK_COMMANDS}")
 fi
 
 add-zsh-hook preexec __savvy_record_pre_exec__
diff --git a/shell/zsh.go b/shell/zsh.go
index 5ecf6ea..e166ed3 100644
--- a/shell/zsh.go
+++ b/shell/zsh.go
@@ -9,6 +9,7 @@ import (
 	"os/user"
 	"path/filepath"
 	"slices"
+	"strconv"
 	"strings"
 	"text/template"
 	"time"
@@ -249,8 +250,14 @@ func (z *zsh) SpawnRunbookRunner(ctx context.Context, runbook *client.Runbook) (
 		nextStep = "1"
 	}
 
+	var nextStepIdx int
+	nextStepIdx, err = strconv.Atoi(nextStep)
+	if err != nil {
+		nextStepIdx = 1
+	}
+
 	cmd := exec.CommandContext(ctx, z.shellCmd)
-	cmd.Env = append(os.Environ(), "ZDOTDIR="+tmp, "SAVVY_CONTEXT=run", fmt.Sprintf("SAVVY_RUNBOOK_COMMANDS=%s", strings.Join(runbook.Commands(), ",")), fmt.Sprintf("SAVVY_NEXT_STEP=%s", nextStep))
+	cmd.Env = append(os.Environ(), "ZDOTDIR="+tmp, "SAVVY_CONTEXT=run", fmt.Sprintf("SAVVY_RUNBOOK_COMMANDS=%s", strings.Join(runbook.Commands(), ",")), fmt.Sprintf("SAVVY_NEXT_STEP=%d", nextStepIdx))
 	cmd.WaitDelay = 500 * time.Millisecond
 	return cmd, nil
 }