fix(mail): replace os.Exit with graceful shutdown in mail watch

[cookier]

Replace os.Exit(0) in the signal handler of mail +watch with context-based graceful shutdown using context.WithCancel and a shutdown channel.

The os.Exit(0) call had several problems:

• Bypasses deferred cleanup functions in the call stack
• Makes the code untestable (process terminates immediately)
• Not idiomatic for Go CLI tools

Changes:

• Use context.WithCancel to create a cancellable watch context
• Signal handler now cancels the context and closes shutdownCh
• cli.Start uses watchCtx instead of the parent context
• Distinguish between signal-triggered shutdown and real errors via select on shutdownCh

Closes #268

Summary

Changes

• Change 1
• Change 2

Test Plan

• Unit tests pass
• Manual local verification confirms the lark xxx command works as expected

Related Issues

• None

Summary by CodeRabbit

• Bug Fixes

  • Improved graceful shutdown behavior when the application receives system interrupt signals
  • Enhanced error handling to distinguish between intentional signal-triggered shutdown and genuine connection failures
  • Refined signal handler implementation for cleaner application termination

• Refactor

  • Streamlined shutdown coordination and synchronization mechanism

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

shaoqi seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Introduces a cancellable watchCtx and coordinated shutdown in mail watch: the signal handler stops signal delivery, cancels watchCtx, and triggers unsubscribe without calling os.Exit(0). WebSocket startup uses watchCtx, and startup errors are handled differently when cancellation is observed.

Changes

Cohort / File(s)	Summary
Mail watch implementation shortcuts/mail/mail_watch.go	Added watchCtx via context.WithCancel, switched cli.Start to use watchCtx, stopped calling os.Exit(0) in signal handler, added signal.Stop(sigCh) and coordinated shutdown/unsubscribe logic; introduced handleMailWatchStartError helper to distinguish canceled vs. real startup errors.
Tests for startup error handling shortcuts/mail/mail_watch_test.go	Added unit tests for handleMailWatchStartError: verifies canceled watchCtx returns nil without calling unsubscribe, and non-context errors call unsubscribe and return a wrapped error.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant Main
    participant SignalHandler
    participant CLI as WebSocketClient
    participant Mailbox

    User->>Main: start `mail watch` command (ctx)
    Main->>Main: create watchCtx, shutdownCh
    Main->>SignalHandler: register signal handler (sigCh)
    Main->>CLI: Start(watchCtx)
    CLI->>Mailbox: open websocket subscribe
    SignalHandler->>SignalHandler: receives interrupt (Ctrl+C)
    SignalHandler->>SignalHandler: signal.Stop(sigCh)
    SignalHandler->>Main: cancelWatch() (cancel watchCtx)
    SignalHandler->>Main: close shutdownCh
    WatchCtx-->>CLI: canceled
    CLI->>Main: Start returns error (context canceled)
    Main->>Mailbox: unsubscribe() or skip if canceled
    Main->>Main: graceful shutdown completes

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related issues

• fix: replace os.Exit with graceful shutdown in mail watch #269: Matches the explicit objective to remove os.Exit(0) and implement graceful shutdown with context cancellation and coordinated signaling.
• fix: remove os.Exit from mail watch signal handler #268: Similar scope — both address replacing abrupt exit with coordinated shutdown in mail watch.

Possibly related PRs

• fix(mail): on-demand scope checks and watch event filtering #198: Modifies shortcuts/mail/mail_watch.go to introduce coordinated shutdown/unsubscribe logic and context-aware WebSocket startup; closely related to these changes.

Suggested reviewers

• infeng

Poem

🐰 Soft thump of Ctrl+C in the night,
I cancel the watch instead of taking flight.
Signals hush, websockets close with grace,
The rabbit stays—no abrupt exit trace.
🎉

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	PR description provides clear motivation and key changes but lacks proper alignment with template structure and incomplete test plan details.	Fill template sections correctly: move main description to Summary section, expand Changes list with specific details (context.WithCancel, signal handler logic, handleMailWatchStartError function), and clarify test plan status (unit tests added or pending).

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The pull request title accurately describes the main change: replacing os.Exit with graceful shutdown coordination using context cancellation in the mail watch functionality.
Linked Issues check	✅ Passed	The pull request implements all primary coding requirements from issue #269: replaces os.Exit with graceful shutdown using context.WithCancel, moves signal.Stop to defer, and distinguishes signal-triggered shutdown from errors.
Out of Scope Changes check	✅ Passed	All changes are directly related to the graceful shutdown implementation objective: signal handler refactoring, context cancellation setup, error handling in cli.Start, and corresponding unit tests for the new handleMailWatchStartError function.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@shortcuts/mail/mail_watch.go`:
- Around line 448-449: The shutdown channel is closed after cancelling the
watch, creating a race where cli.Start(watchCtx) may return before shutdownCh is
closed and the subsequent select treats it as a network error; to fix, close
shutdownCh before calling cancelWatch() so the signal is visible to any
goroutine that returns from cli.Start(watchCtx), i.e., swap the order of
close(shutdownCh) and cancelWatch() around the watch shutdown logic while
keeping the same ownership semantics for shutdownCh and cancelWatch().

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 462e0b75-29eb-4ca6-8b15-55e774b8450d

📥 Commits

Reviewing files that changed from the base of the PR and between 0c77c95 and 978edc6.

📒 Files selected for processing (1)

• shortcuts/mail/mail_watch.go

[greptile-apps]

Greptile Summary

Replaces os.Exit(0) in the mail watch signal handler with context-based graceful shutdown, fixing deferred-cleanup bypass and testability issues. The new handleMailWatchStartError helper uses watchCtx.Err() to atomically distinguish signal-triggered cancellation from genuine WebSocket failures, directly addressing the race condition flagged in the previous review.

Confidence Score: 5/5

Safe to merge — the previous P1 race condition is resolved and no new issues were found.

The race between cancelWatch() and a shutdownCh is eliminated by using watchCtx.Err() as recommended by the prior review. Both the graceful-shutdown and real-error paths have unit test coverage. No P0/P1 issues remain.

No files require special attention.

Important Files Changed

Filename	Overview
shortcuts/mail/mail_watch.go	Replaces os.Exit(0) with context cancellation; handleMailWatchStartError uses watchCtx.Err() for race-free shutdown detection
shortcuts/mail/mail_watch_test.go	Adds two unit tests covering graceful shutdown and real network failure paths for handleMailWatchStartError

Sequence Diagram

sequenceDiagram
    participant Main
    participant SignalHandler
    participant WebSocket as cli.Start

    Main->>WebSocket: cli.Start(watchCtx)
    Note over Main: Blocks — waiting for mail events

    SignalHandler-->>SignalHandler: Receives SIGINT/SIGTERM
    SignalHandler->>SignalHandler: signal.Stop(sigCh)
    SignalHandler->>SignalHandler: unsubscribe()
    SignalHandler->>Main: cancelWatch()

    WebSocket-->>Main: returns error (context canceled)
    Main->>Main: handleMailWatchStartError(err, watchCtx, ...)
    Note over Main: watchCtx.Err() != nil → return nil

Loading

_{Greploops — Automatically fix all review issues by running /greploops in Claude Code. It iterates: fix, push, re-review, repeat until 5/5 confidence.
Use the Greptile plugin for Claude Code to query reviews, search comments, and manage custom context directly from your terminal.}

_{Reviews (2): Last reviewed commit: "fix(mail): avoid shutdown race in mail w..." | Re-trigger Greptile}

[cookier]

Fixed. I removed the shutdownCh-based classification and now use watchCtx.Err() after cli.Start(watchCtx) returns.

That removes the race between cancelWatch() and channel close, avoids reporting graceful shutdown as a network error, and prevents the extra cleanup path from running on signal-triggered exit.

[haidaodashushu]

Found one cleanup edge case in the new graceful-shutdown path.

[haidaodashushu]

watchCtx.Err() != nil is broader than "graceful shutdown via SIGINT/SIGTERM" here. watchCtx is derived from the parent ctx, so external cancellation from the caller/runtime will also hit this branch and return nil immediately, even though only the signal path actually performs unsubscribe() before canceling the watch. That means non-signal context cancellation can now skip mailbox cleanup entirely. Please distinguish signal-triggered shutdown from generic context cancellation instead of using watchCtx.Err() != nil as the only condition.