diff --git a/shortcuts/mail/mail_watch.go b/shortcuts/mail/mail_watch.go index c5699427..0c1d84a2 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" @@ -259,19 +260,30 @@ 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 } - eventCount := 0 + var eventCount int64 handleEvent := func(data map[string]interface{}) { // Extract event body @@ -337,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. @@ -426,27 +438,41 @@ 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 { 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: + // 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(): + return } - signal.Stop(sigCh) - os.Exit(0) }() info("Connected. Waiting for mail events... (Ctrl+C to stop)") - if err := cli.Start(ctx); err != nil { - unsubscribe() //nolint:errcheck // best-effort cleanup + if err := cli.Start(watchCtx); err != nil { + select { + case <-shutdownBySignal: + return nil + default: + } + if errors.Is(err, context.Canceled) { + return nil + } return output.ErrNetwork("WebSocket connection failed: %v", err) } return nil