From 7e45bcad6670155fd730c4ad7d542e76783dd177 Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Sun, 5 Apr 2026 17:19:44 +0800
Subject: [PATCH 01/10] fix: replace os.Exit with graceful shutdown in mail
 watch

Remove os.Exit(0) from signal handler and implement proper shutdown
coordination using context cancellation and channel signaling. This
makes the code testable and allows proper cleanup of resources.
---
 shortcuts/mail/mail_watch.go | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index c56994270..d4fd47e5e 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -426,6 +426,12 @@ var MailWatch = common.Shortcut{
 
 		sigCh := make(chan os.Signal, 1)
 		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+		defer signal.Stop(sigCh)
+
+		watchCtx, cancelWatch := context.WithCancel(ctx)
+		defer cancelWatch()
+
+		shutdownBySignal := make(chan struct{})
 		go func() {
 			defer func() {
 				if r := recover(); r != nil {
@@ -440,12 +446,21 @@ var MailWatch = common.Shortcut{
 			} else {
 				info("Mailbox unsubscribed.")
 			}
-			signal.Stop(sigCh)
-			os.Exit(0)
+			cancelWatch()
+			close(shutdownBySignal)
 		}()
 
 		info("Connected. Waiting for mail events... (Ctrl+C to stop)")
-		if err := cli.Start(ctx); err != nil {
+		if err := cli.Start(watchCtx); err != nil {
+			select {
+			case <-shutdownBySignal:
+				return nil
+			default:
+			}
+			if errors.Is(err, context.Canceled) {
+				unsubscribe() //nolint:errcheck // best-effort cleanup
+				return nil
+			}
 			unsubscribe() //nolint:errcheck // best-effort cleanup
 			return output.ErrNetwork("WebSocket connection failed: %v", err)
 		}

From 9c934018c620e703c66ced3e38f304ee6d25ff3a Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Sun, 5 Apr 2026 17:44:06 +0800
Subject: [PATCH 02/10] fix: order signal shutdown markers in mail watch

---
 shortcuts/mail/mail_watch.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index d4fd47e5e..594cbf5bb 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -446,8 +446,8 @@ var MailWatch = common.Shortcut{
 			} else {
 				info("Mailbox unsubscribed.")
 			}
-			cancelWatch()
 			close(shutdownBySignal)
+			cancelWatch()
 		}()
 
 		info("Connected. Waiting for mail events... (Ctrl+C to stop)")

From c5e7b5a3a8e48a08f724ae2afb7e93e246ce4a52 Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Sun, 5 Apr 2026 18:00:40 +0800
Subject: [PATCH 03/10] fix: improve signal handler to avoid blocking shutdown

- Add select with watchCtx.Done() to handle non-signal exit paths
- Move close(shutdownBySignal) and cancelWatch() before unsubscribe()
- Execute unsubscribe() in main goroutine after cli.Start() returns

This ensures:
- Signal handler goroutine can exit on non-signal paths
- Shutdown signal is published immediately without blocking on network I/O
- No race condition between shutdownBySignal and context.Canceled
---
 shortcuts/mail/mail_watch.go | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index 594cbf5bb..cc663509a 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -438,27 +438,36 @@ var MailWatch = common.Shortcut{
 					fmt.Fprintf(errOut, "panic in signal handler: %v\n", r)
 				}
 			}()
-			<-sigCh
-			info(fmt.Sprintf("\nShutting down... (received %d events)", eventCount))
-			info("Unsubscribing mailbox events...")
-			if unsubErr := unsubscribe(); unsubErr != nil {
-				fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", unsubErr)
-			} else {
-				info("Mailbox unsubscribed.")
+			select {
+			case <-sigCh:
+				info(fmt.Sprintf("\nShutting down... (received %d events)", eventCount))
+				close(shutdownBySignal)
+				cancelWatch()
+			case <-watchCtx.Done():
+				return
 			}
-			close(shutdownBySignal)
-			cancelWatch()
 		}()
 
 		info("Connected. Waiting for mail events... (Ctrl+C to stop)")
 		if err := cli.Start(watchCtx); err != nil {
 			select {
 			case <-shutdownBySignal:
+				info("Unsubscribing mailbox events...")
+				if unsubErr := unsubscribe(); unsubErr != nil {
+					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", unsubErr)
+				} else {
+					info("Mailbox unsubscribed.")
+				}
 				return nil
 			default:
 			}
 			if errors.Is(err, context.Canceled) {
-				unsubscribe() //nolint:errcheck // best-effort cleanup
+				info("Unsubscribing mailbox events...")
+				if unsubErr := unsubscribe(); unsubErr != nil {
+					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", unsubErr)
+				} else {
+					info("Mailbox unsubscribed.")
+				}
 				return nil
 			}
 			unsubscribe() //nolint:errcheck // best-effort cleanup

From e4656a4b3c2b185c84796908bcce37fa70863000 Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Sun, 5 Apr 2026 18:14:15 +0800
Subject: [PATCH 04/10] refactor: use defer + sync.Once for cleanup on all exit
 paths

- Create unsubscribeWithLog() with sync.Once to ensure single execution
- Use defer to guarantee cleanup on all exit paths
- Remove duplicate unsubscribe logic from error handling branches

This ensures unsubscribe runs exactly once regardless of how the function exits:
- Normal exit
- Early return (e.g., profile error)
- Signal-triggered shutdown
- Context cancellation
- Network errors
---
 shortcuts/mail/mail_watch.go | 26 ++++++++++++--------------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index cc663509a..33c8fb136 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -259,13 +259,24 @@ var MailWatch = common.Shortcut{
 			})
 			return unsubErr
 		}
+		var unsubscribeLogOnce sync.Once
+		unsubscribeWithLog := func() {
+			unsubscribeLogOnce.Do(func() {
+				info("Unsubscribing mailbox events...")
+				if err := unsubscribe(); err != nil {
+					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", err)
+				} else {
+					info("Mailbox unsubscribed.")
+				}
+			})
+		}
+		defer unsubscribeWithLog()
 
 		// Resolve "me" to the actual email address so we can filter events.
 		mailboxFilter := mailbox
 		if mailbox == "me" {
 			resolved, profileErr := fetchMailboxPrimaryEmail(runtime, "me")
 			if profileErr != nil {
-				unsubscribe() //nolint:errcheck // best-effort cleanup; primary error is profileErr
 				return enhanceProfileError(profileErr)
 			}
 			mailboxFilter = resolved
@@ -452,25 +463,12 @@ var MailWatch = common.Shortcut{
 		if err := cli.Start(watchCtx); err != nil {
 			select {
 			case <-shutdownBySignal:
-				info("Unsubscribing mailbox events...")
-				if unsubErr := unsubscribe(); unsubErr != nil {
-					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", unsubErr)
-				} else {
-					info("Mailbox unsubscribed.")
-				}
 				return nil
 			default:
 			}
 			if errors.Is(err, context.Canceled) {
-				info("Unsubscribing mailbox events...")
-				if unsubErr := unsubscribe(); unsubErr != nil {
-					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", unsubErr)
-				} else {
-					info("Mailbox unsubscribed.")
-				}
 				return nil
 			}
-			unsubscribe() //nolint:errcheck // best-effort cleanup
 			return output.ErrNetwork("WebSocket connection failed: %v", err)
 		}
 		return nil

From 8ea53f9ea0cc834ac3345b550728b33145d2266a Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Sun, 5 Apr 2026 18:30:38 +0800
Subject: [PATCH 05/10] fix: harden mail watch signal shutdown path

---
 shortcuts/mail/mail_watch.go | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index 33c8fb136..0c1d84a2b 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -18,6 +18,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 
 	larkcore "github.com/larksuite/oapi-sdk-go/v3/core"
@@ -282,7 +283,7 @@ var MailWatch = common.Shortcut{
 			mailboxFilter = resolved
 		}
 
-		eventCount := 0
+		var eventCount int64
 
 		handleEvent := func(data map[string]interface{}) {
 			// Extract event body
@@ -348,7 +349,7 @@ var MailWatch = common.Shortcut{
 				}
 			}
 
-			eventCount++
+			atomic.AddInt64(&eventCount, 1)
 
 			// Prompt injection detection: warn when email body contains known injection patterns.
 			// Body fields may be base64url-encoded; decode before scanning.
@@ -451,7 +452,10 @@ var MailWatch = common.Shortcut{
 			}()
 			select {
 			case <-sigCh:
-				info(fmt.Sprintf("\nShutting down... (received %d events)", eventCount))
+				// Restore default signal behavior so a second Ctrl+C can force terminate.
+				signal.Stop(sigCh)
+				signal.Reset(os.Interrupt, syscall.SIGTERM)
+				info(fmt.Sprintf("\nShutting down... (received %d events)", atomic.LoadInt64(&eventCount)))
 				close(shutdownBySignal)
 				cancelWatch()
 			case <-watchCtx.Done():

From 161916a3c91b48f6a52c79c623c6c713bd67138f Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Wed, 8 Apr 2026 23:39:08 +0800
Subject: [PATCH 06/10] fix(mail): resolve graceful shutdown blocking in
 waitForMailWatchShutdown

Add goroutine with buffered error channel to prevent cli.Start from blocking
on context cancellation. This addresses the review comment in PR #269
where the watch command would hang on shutdown.

Changes:
- Move cli.Start to goroutine with buffered error channel
- Add waitForMailWatchShutdown helper function
- Handle shutdownBySignal immediately without waiting for Start result
- Return nil if context is canceled or error occurs after shutdown
---
 shortcuts/mail/mail_watch.go | 31 ++++++++++++++++++++++---------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index 0c1d84a2b..12c799673 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -77,6 +77,23 @@ func detectPromptInjection(content string) bool {
 	return false
 }
 
+func waitForMailWatchShutdown(startErrCh <-chan error, shutdownBySignal <-chan struct{}) error {
+	select {
+	case <-shutdownBySignal:
+		return nil
+	case err := <-startErrCh:
+		select {
+		case <-shutdownBySignal:
+			return nil
+		default:
+		}
+		if err == nil || errors.Is(err, context.Canceled) {
+			return nil
+		}
+		return err
+	}
+}
+
 var MailWatch = common.Shortcut{
 	Service:     "mail",
 	Command:     "+watch",
@@ -464,15 +481,11 @@ var MailWatch = common.Shortcut{
 		}()
 
 		info("Connected. Waiting for mail events... (Ctrl+C to stop)")
-		if err := cli.Start(watchCtx); err != nil {
-			select {
-			case <-shutdownBySignal:
-				return nil
-			default:
-			}
-			if errors.Is(err, context.Canceled) {
-				return nil
-			}
+		startErrCh := make(chan error, 1)
+		go func() {
+			startErrCh <- cli.Start(watchCtx)
+		}()
+		if err := waitForMailWatchShutdown(startErrCh, shutdownBySignal); err != nil {
 			return output.ErrNetwork("WebSocket connection failed: %v", err)
 		}
 		return nil

From d06bbe1de9edebd5d90458ab2e510a08f4e224e5 Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Wed, 8 Apr 2026 23:39:12 +0800
Subject: [PATCH 07/10] test(mail): add comprehensive tests for
 waitForMailWatchShutdown

Test scenarios for the graceful shutdown helper function:
- Return immediately on shutdownBySignal
- Return nil on context.Canceled
- Propagate other errors when shutdown not signaled
- Ensure non-blocking behavior on signal
---
 shortcuts/mail/mail_watch_test.go | 42 +++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/shortcuts/mail/mail_watch_test.go b/shortcuts/mail/mail_watch_test.go
index 02476fbdf..b73a88378 100644
--- a/shortcuts/mail/mail_watch_test.go
+++ b/shortcuts/mail/mail_watch_test.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/larksuite/cli/internal/core"
 	"github.com/larksuite/cli/internal/output"
@@ -264,6 +265,47 @@ func TestMailWatchLoggerSuppressesDebugAlways(t *testing.T) {
 	}
 }
 
+func TestWaitForMailWatchShutdownReturnsOnSignalWithoutStartResult(t *testing.T) {
+	startErrCh := make(chan error)
+	shutdownBySignal := make(chan struct{})
+	close(shutdownBySignal)
+
+	done := make(chan error, 1)
+	go func() {
+		done <- waitForMailWatchShutdown(startErrCh, shutdownBySignal)
+	}()
+
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("expected nil on signal shutdown, got %v", err)
+		}
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("waitForMailWatchShutdown blocked after signal shutdown")
+	}
+}
+
+func TestWaitForMailWatchShutdownReturnsNilOnContextCanceled(t *testing.T) {
+	startErrCh := make(chan error, 1)
+	shutdownBySignal := make(chan struct{})
+	startErrCh <- context.Canceled
+
+	if err := waitForMailWatchShutdown(startErrCh, shutdownBySignal); err != nil {
+		t.Fatalf("expected nil for context.Canceled, got %v", err)
+	}
+}
+
+func TestWaitForMailWatchShutdownReturnsStartError(t *testing.T) {
+	startErrCh := make(chan error, 1)
+	shutdownBySignal := make(chan struct{})
+	want := assertErr("boom")
+	startErrCh <- want
+
+	if err := waitForMailWatchShutdown(startErrCh, shutdownBySignal); err != want {
+		t.Fatalf("expected original error, got %v", err)
+	}
+}
+
 func TestDecodeBodyFieldsForFileDecodesMessageWrapper(t *testing.T) {
 	htmlEncoded := base64.URLEncoding.EncodeToString([]byte("<h1>Hello</h1>"))
 	plainEncoded := base64.URLEncoding.EncodeToString([]byte("Hello plain"))

From 338c55ac42678e3427154dca9fd9e46db01deeae Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Wed, 8 Apr 2026 23:43:04 +0800
Subject: [PATCH 08/10] fix(mail): propagate unsubscribe failures on graceful
 shutdown paths

- Return non-zero exit code when unsubscribe fails instead of warning
- Ensure proper cleanup before function returns
- Fix issue where command exits 0 even when subscription not cleaned up
---
 shortcuts/mail/mail_watch.go | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index 12c799673..a2412243f 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -278,7 +278,7 @@ var MailWatch = common.Shortcut{
 			return unsubErr
 		}
 		var unsubscribeLogOnce sync.Once
-		unsubscribeWithLog := func() {
+		unsubscribeWithLog := func() error {
 			unsubscribeLogOnce.Do(func() {
 				info("Unsubscribing mailbox events...")
 				if err := unsubscribe(); err != nil {
@@ -287,9 +287,8 @@ var MailWatch = common.Shortcut{
 					info("Mailbox unsubscribed.")
 				}
 			})
+			return unsubErr
 		}
-		defer unsubscribeWithLog()
-
 		// Resolve "me" to the actual email address so we can filter events.
 		mailboxFilter := mailbox
 		if mailbox == "me" {
@@ -488,6 +487,12 @@ var MailWatch = common.Shortcut{
 		if err := waitForMailWatchShutdown(startErrCh, shutdownBySignal); err != nil {
 			return output.ErrNetwork("WebSocket connection failed: %v", err)
 		}
+
+		// Ensure cleanup happens before returning
+		unsubErr = unsubscribeWithLog()
+		if unsubErr != nil {
+			return output.ErrSystem("Failed to unsubscribe mailbox events: %v", unsubErr)
+		}
 		return nil
 	},
 }

From 3c93d5c6aeca79b28007259bfab1b17ac4ff758b Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Wed, 8 Apr 2026 23:50:24 +0800
Subject: [PATCH 09/10] fix(mail): restore unified watch cleanup propagation

---
 shortcuts/mail/mail_watch.go | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go
index a2412243f..a318b1990 100644
--- a/shortcuts/mail/mail_watch.go
+++ b/shortcuts/mail/mail_watch.go
@@ -94,6 +94,16 @@ func waitForMailWatchShutdown(startErrCh <-chan error, shutdownBySignal <-chan s
 	}
 }
 
+func finalizeMailWatchCleanup(runErr, cleanupErr error) error {
+	if cleanupErr == nil {
+		return runErr
+	}
+	if runErr != nil {
+		return runErr
+	}
+	return cleanupErr
+}
+
 var MailWatch = common.Shortcut{
 	Service:     "mail",
 	Command:     "+watch",
@@ -180,7 +190,7 @@ var MailWatch = common.Shortcut{
 		}
 		return d
 	},
-	Execute: func(ctx context.Context, runtime *common.RuntimeContext) error {
+	Execute: func(ctx context.Context, runtime *common.RuntimeContext) (retErr error) {
 		if runtime.Bool("print-output-schema") {
 			printWatchOutputSchema(runtime)
 			return nil
@@ -278,17 +288,22 @@ var MailWatch = common.Shortcut{
 			return unsubErr
 		}
 		var unsubscribeLogOnce sync.Once
-		unsubscribeWithLog := func() error {
+		unsubscribeWithLog := func(primaryErr error) error {
 			unsubscribeLogOnce.Do(func() {
 				info("Unsubscribing mailbox events...")
 				if err := unsubscribe(); err != nil {
-					fmt.Fprintf(errOut, "Warning: unsubscribe failed: %v\n", err)
+					if primaryErr != nil {
+						fmt.Fprintf(errOut, "Warning: unsubscribe failed during cleanup: %v\n", err)
+					}
 				} else {
 					info("Mailbox unsubscribed.")
 				}
 			})
 			return unsubErr
 		}
+		defer func() {
+			retErr = finalizeMailWatchCleanup(retErr, unsubscribeWithLog(retErr))
+		}()
 		// Resolve "me" to the actual email address so we can filter events.
 		mailboxFilter := mailbox
 		if mailbox == "me" {
@@ -487,12 +502,6 @@ var MailWatch = common.Shortcut{
 		if err := waitForMailWatchShutdown(startErrCh, shutdownBySignal); err != nil {
 			return output.ErrNetwork("WebSocket connection failed: %v", err)
 		}
-
-		// Ensure cleanup happens before returning
-		unsubErr = unsubscribeWithLog()
-		if unsubErr != nil {
-			return output.ErrSystem("Failed to unsubscribe mailbox events: %v", unsubErr)
-		}
 		return nil
 	},
 }

From c258e9d3be686d8bc63e550fb3f9cad5e24aad5a Mon Sep 17 00:00:00 2001
From: luojiyin <luojiyin@hotmail.com>
Date: Wed, 8 Apr 2026 23:50:39 +0800
Subject: [PATCH 10/10] test(mail): cover unified watch cleanup result
 selection

---
 shortcuts/mail/mail_watch_test.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/shortcuts/mail/mail_watch_test.go b/shortcuts/mail/mail_watch_test.go
index b73a88378..8139b3c22 100644
--- a/shortcuts/mail/mail_watch_test.go
+++ b/shortcuts/mail/mail_watch_test.go
@@ -306,6 +306,24 @@ func TestWaitForMailWatchShutdownReturnsStartError(t *testing.T) {
 	}
 }
 
+func TestFinalizeMailWatchCleanup(t *testing.T) {
+	runErr := assertErr("run failed")
+	cleanupErr := assertErr("cleanup failed")
+
+	if err := finalizeMailWatchCleanup(nil, nil); err != nil {
+		t.Fatalf("expected nil, got %v", err)
+	}
+	if err := finalizeMailWatchCleanup(runErr, nil); err != runErr {
+		t.Fatalf("expected runErr when cleanup succeeds, got %v", err)
+	}
+	if err := finalizeMailWatchCleanup(nil, cleanupErr); err != cleanupErr {
+		t.Fatalf("expected cleanupErr when run succeeds, got %v", err)
+	}
+	if err := finalizeMailWatchCleanup(runErr, cleanupErr); err != runErr {
+		t.Fatalf("expected primary runErr to win, got %v", err)
+	}
+}
+
 func TestDecodeBodyFieldsForFileDecodesMessageWrapper(t *testing.T) {
 	htmlEncoded := base64.URLEncoding.EncodeToString([]byte("<h1>Hello</h1>"))
 	plainEncoded := base64.URLEncoding.EncodeToString([]byte("Hello plain"))