From ba96fa06feeae25d855886b91cff868191b1bce4 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 30 Sep 2025 11:20:37 +0530
Subject: [PATCH] Allow using defer in KittenMain

---
 kittens/ssh/askpass.go | 29 +++++++++++++++--------------
 tools/cli/command.go   |  7 -------
 tools/cmd/main.go      |  9 ++++-----
 3 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/kittens/ssh/askpass.go b/kittens/ssh/askpass.go
index bf086ffca..854d84399 100644
--- a/kittens/ssh/askpass.go
+++ b/kittens/ssh/askpass.go
@@ -16,25 +16,25 @@ import (
 
 var _ = fmt.Print
 
-func fatal(err error) {
+func fatal(err error) int {
 	cli.ShowError(err)
-	os.Exit(1)
+	return 1
 }
 
-func trigger_ask(name string) {
+func trigger_ask(name string) int {
 	term, err := tty.OpenControllingTerm()
 	if err != nil {
-		fatal(err)
+		return fatal(err)
 	}
 	defer term.Close()
 	_, err = term.WriteString("\x1bP@kitty-ask|" + name + "\x1b\\")
 	if err != nil {
-		fatal(err)
+		return fatal(err)
 	}
-
+	return 0
 }
 
-func RunSSHAskpass() {
+func RunSSHAskpass() int {
 	msg := os.Args[len(os.Args)-1]
 	prompt := os.Getenv("SSH_ASKPASS_PROMPT")
 	is_confirm := prompt == "confirm"
@@ -50,21 +50,21 @@ func RunSSHAskpass() {
 	}
 	data, err := json.Marshal(q)
 	if err != nil {
-		fatal(err)
+		return fatal(err)
 	}
 	data_shm, err := shm.CreateTemp("askpass-*", uint64(len(data)+32))
 	if err != nil {
-		fatal(fmt.Errorf("Failed to create SHM file with error: %w", err))
+		return fatal(fmt.Errorf("Failed to create SHM file with error: %w", err))
 	}
 	defer data_shm.Close()
 	defer func() { _ = data_shm.Unlink() }()
 
 	data_shm.Slice()[0] = 0
 	if err = shm.WriteWithSize(data_shm, data, 1); err != nil {
-		fatal(fmt.Errorf("Failed to write to SHM file with error: %w", err))
+		return fatal(fmt.Errorf("Failed to write to SHM file with error: %w", err))
 	}
 	if err = data_shm.Flush(); err != nil {
-		fatal(fmt.Errorf("Failed to flush SHM file with error: %w", err))
+		return fatal(fmt.Errorf("Failed to flush SHM file with error: %w", err))
 	}
 	trigger_ask(data_shm.Name())
 	for {
@@ -75,14 +75,14 @@ func RunSSHAskpass() {
 	}
 	data, err = shm.ReadWithSize(data_shm, 1)
 	if err != nil {
-		fatal(fmt.Errorf("Failed to read from SHM file with error: %w", err))
+		return fatal(fmt.Errorf("Failed to read from SHM file with error: %w", err))
 	}
 	response := ""
 	if is_confirm {
 		var ok bool
 		err = json.Unmarshal(data, &ok)
 		if err != nil {
-			fatal(fmt.Errorf("Failed to parse response data: %#v with error: %w", string(data), err))
+			return fatal(fmt.Errorf("Failed to parse response data: %#v with error: %w", string(data), err))
 		}
 		response = "no"
 		if ok {
@@ -91,7 +91,7 @@ func RunSSHAskpass() {
 	} else {
 		err = json.Unmarshal(data, &response)
 		if err != nil {
-			fatal(fmt.Errorf("Failed to parse response data: %#v with error: %w", string(data), err))
+			return fatal(fmt.Errorf("Failed to parse response data: %#v with error: %w", string(data), err))
 		}
 		if is_fingerprint_check {
 			response = strings.ToLower(response)
@@ -106,4 +106,5 @@ func RunSSHAskpass() {
 	if response != "" {
 		fmt.Println(response)
 	}
+	return 0
 }
diff --git a/tools/cli/command.go b/tools/cli/command.go
index a1d74d029..974cbc76d 100644
--- a/tools/cli/command.go
+++ b/tools/cli/command.go
@@ -565,13 +565,6 @@ func (self *Command) ExecArgs(args []string) (exit_code int) {
 	return
 }
 
-func (self *Command) Exec(args ...string) {
-	if len(args) == 0 {
-		args = os.Args
-	}
-	os.Exit(self.ExecArgs(args))
-}
-
 func (self *Command) GetCompletions(argv []string, init_completions func(*Completions)) *Completions {
 	ans := NewCompletions()
 	if init_completions != nil {
diff --git a/tools/cmd/main.go b/tools/cmd/main.go
index 1c8aa0c1c..bcbf2efa5 100644
--- a/tools/cmd/main.go
+++ b/tools/cmd/main.go
@@ -16,13 +16,12 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func KittenMain(args ...string) {
+func KittenMain(args ...string) int {
 	krm := os.Getenv("KITTY_KITTEN_RUN_MODULE")
 	os.Unsetenv("KITTY_KITTEN_RUN_MODULE")
 	switch krm {
 	case "ssh_askpass":
-		ssh.RunSSHAskpass()
-		return
+		return ssh.RunSSHAskpass()
 	}
 	root := cli.NewRootCommand()
 	root.ShortDescription = "Fast, statically compiled implementations of various kittens (command line tools for use with kitty)"
@@ -49,9 +48,9 @@ func KittenMain(args ...string) {
 	completion.EntryPoint(root)
 
 	root.SubCommandIsOptional = true
-	root.Exec(args...)
+	return root.ExecArgs(args)
 }
 
 func main() {
-	KittenMain()
+	os.Exit(KittenMain())
 }