From c26b39a9aa56e834262753e08372240fde9dcdf1 Mon Sep 17 00:00:00 2001
From: zeertzjq <zeertzjq@outlook.com>
Date: Sat, 23 Sep 2023 11:00:20 +0800
Subject: [PATCH] fix(tui): make :cquit work properly with remote TUI (#25313)

---
 src/nvim/main.c                       |  3 ++
 src/nvim/tui/tui.c                    |  3 ++
 test/functional/terminal/tui_spec.lua | 66 ++++++++++++++++++++++++---
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/src/nvim/main.c b/src/nvim/main.c
index d9ca82784f..0bb69d000b 100644
--- a/src/nvim/main.c
+++ b/src/nvim/main.c
@@ -659,6 +659,9 @@ void os_exit(int r)
 
   if (ui_client_channel_id) {
     ui_client_stop();
+    if (r == 0) {
+      r = ui_client_exit_status;
+    }
   } else {
     ui_flush();
     ui_call_stop();
diff --git a/src/nvim/tui/tui.c b/src/nvim/tui/tui.c
index 4097b770c9..db15cdb053 100644
--- a/src/nvim/tui/tui.c
+++ b/src/nvim/tui/tui.c
@@ -386,6 +386,9 @@ static void terminfo_stop(TUIData *tui)
   unibi_out_ext(tui, tui->unibi_ext.disable_extended_keys);
   // May restore old title before exiting alternate screen.
   tui_set_title(tui, (String)STRING_INIT);
+  if (ui_client_exit_status == 0) {
+    ui_client_exit_status = tui->seen_error_exit;
+  }
   // if nvim exited with nonzero status, without indicated this was an
   // intentional exit (like `:1cquit`), it likely was an internal failure.
   // Don't clobber the stderr error message in this case.
diff --git a/test/functional/terminal/tui_spec.lua b/test/functional/terminal/tui_spec.lua
index 2d4613dda4..f8d2c01771 100644
--- a/test/functional/terminal/tui_spec.lua
+++ b/test/functional/terminal/tui_spec.lua
@@ -1720,7 +1720,15 @@ describe('TUI', function()
 
   it('no assert failure on deadly signal #21896', function()
     exec_lua([[vim.uv.kill(vim.fn.jobpid(vim.bo.channel), 'sigterm')]])
-    screen:expect({any = '%[Process exited 1%]'})
+    screen:expect{grid=[[
+      Vim: Caught deadly signal 'SIGTERM'               |
+                                                        |
+                                                        |
+      [Process exited 1]{1: }                               |
+                                                        |
+                                                        |
+      {3:-- TERMINAL --}                                    |
+    ]]}
   end)
 
   it('no stack-use-after-scope with cursor color #22432', function()
@@ -2743,7 +2751,15 @@ describe("TUI as a client", function()
     -- No heap-use-after-free when receiving UI events after deadly signal #22184
     server:request('nvim_input', ('a'):rep(1000))
     exec_lua([[vim.uv.kill(vim.fn.jobpid(vim.bo.channel), 'sigterm')]])
-    screen:expect({any = '%[Process exited 1%]'})
+    screen:expect{grid=[[
+      Vim: Caught deadly signal 'SIGTERM'               |
+                                                        |
+                                                        |
+      [Process exited 1]{1: }                               |
+                                                        |
+                                                        |
+      {3:-- TERMINAL --}                                    |
+    ]]}
 
     eq(0, meths.get_vvar('shell_error'))
     -- exits on input eof #22244
@@ -2771,7 +2787,7 @@ describe("TUI as a client", function()
     ]])
   end)
 
-  it("exits when server quits", function()
+  local function test_remote_tui_quit(status)
     local server_super = spawn_argv(false) -- equivalent to clear()
     local client_super = spawn_argv(true)
 
@@ -2780,6 +2796,15 @@ describe("TUI as a client", function()
     local screen_server = thelpers.screen_setup(0,
       string.format([=[["%s", "--listen", "%s", "-u", "NONE", "-i", "NONE", "--cmd", "%s laststatus=2 background=dark"]]=],
         nvim_prog, server_pipe, nvim_set))
+    screen_server:expect{grid=[[
+      {1: }                                                 |
+      {4:~                                                 }|
+      {4:~                                                 }|
+      {4:~                                                 }|
+      {5:[No Name]                                         }|
+                                                        |
+      {3:-- TERMINAL --}                                    |
+    ]]}
 
     feed_data("iHello, World")
     screen_server:expect{grid=[[
@@ -2819,13 +2844,40 @@ describe("TUI as a client", function()
 
     -- quitting the server
     set_session(server_super)
-    feed_data(":q!\n")
-    screen_server:expect({any="Process exited 0"})
-
+    feed_data(status and ':' .. status .. 'cquit!\n' or ":quit!\n")
+    status = status and status or 0
+    screen_server:expect{grid=[[
+                                                        |
+      [Process exited ]] .. status .. [[]{1: }{MATCH:%s+}|
+                                                        |
+                                                        |
+                                                        |
+                                                        |
+      {3:-- TERMINAL --}                                    |
+    ]]}
     -- assert that client has exited
-    screen_client:expect({any="Process exited 0"})
+    screen_client:expect{grid=[[
+                                                        |
+      [Process exited ]] .. status .. [[]{1: }{MATCH:%s+}|
+                                                        |
+                                                        |
+                                                        |
+                                                        |
+      {3:-- TERMINAL --}                                    |
+    ]]}
+
 
     server_super:close()
     client_super:close()
+  end
+
+  describe("exits when server quits", function()
+    it("with :quit", function()
+      test_remote_tui_quit()
+    end)
+
+    it("with :cquit", function()
+      test_remote_tui_quit(42)
+    end)
   end)
 end)