1 files changed, 24 insertions, 1 deletions
diff --git a/src/nvim/msgpack_rpc/channel.c b/src/nvim/msgpack_rpc/channel.c
index 0c04a7b23e..4c35cce09a 100644
--- a/src/nvim/msgpack_rpc/channel.c
+++ b/src/nvim/msgpack_rpc/channel.c
@@ -10,6 +10,7 @@
 #include "nvim/api/private/helpers.h"
 #include "nvim/api/vim.h"
 #include "nvim/msgpack_rpc/channel.h"
+#include "nvim/msgpack_rpc/remote_ui.h"
 #include "nvim/os/event.h"
 #include "nvim/os/rstream.h"
 #include "nvim/os/rstream_defs.h"
@@ -100,6 +101,17 @@ void channel_init(void)
   if (embedded_mode) {
     channel_from_stdio();
   }
+
+  if (abstract_ui) {
+    // Add handler for "attach_ui"
+    remote_ui_init();
+    String method = cstr_as_string("attach_ui");
+    MsgpackRpcRequestHandler handler = {.fn = remote_ui_attach, .defer = true};
+    msgpack_rpc_add_method_handler(method, handler);
+    method = cstr_as_string("detach_ui");
+    handler.fn = remote_ui_detach;
+    msgpack_rpc_add_method_handler(method, handler);
+  }
 }
 
 /// Teardown the module
@@ -241,6 +253,7 @@ Object channel_send_call(uint64_t id,
 
   if (frame.errored) {
     api_set_error(err, Exception, "%s", frame.result.data.string.data);
+    api_free_object(frame.result);
     return NIL;
   }
 
@@ -347,7 +360,13 @@ static void job_err(RStream *rstream, void *data, bool eof)
 
 static void job_exit(Job *job, void *data)
 {
-  free_channel((Channel *)data);
+  Channel *channel = data;
+  // ensure the channel is flagged as closed so channel_send_call frees it
+  // later
+  channel->closed = true;
+  if (!kv_size(channel->call_stack)) {
+    free_channel(channel);
+  }
 }
 
 static void parse_msgpack(RStream *rstream, void *data, bool eof)
@@ -638,6 +657,10 @@ static void on_stdio_close(Event e)
 
 static void free_channel(Channel *channel)
 {
+  if (abstract_ui) {
+    remote_ui_disconnect(channel->id);
+  }
+
   pmap_del(uint64_t)(channels, channel->id);
   msgpack_unpacker_free(channel->unpacker);