Commit 8797d71

mo khan <mo@mokhan.ca>
2025-08-14 23:08:42
feat: perform graceful mcp client shutdown
1 parent 3c5caa4
lib/elelem/agent.rb
@@ -11,6 +11,9 @@ module Elelem
       @model = configuration.model
       @conversation = configuration.conversation
       @logger = configuration.logger
+
+      at_exit { cleanup }
+
       transition_to(States::Idle.new)
     end
 
@@ -32,9 +35,15 @@ module Elelem
 
     def quit
       logger.debug("Exiting...")
+      cleanup
       exit
     end
 
+    def cleanup
+      logger.debug("Cleaning up agent...")
+      configuration.cleanup
+    end
+
     private
 
     attr_reader :configuration, :current_state
lib/elelem/configuration.rb
@@ -44,6 +44,10 @@ module Elelem
       @tools ||= Tools.new(self, [Toolbox::Bash.new(self)] + mcp_tools)
     end
 
+    def cleanup
+      @mcp_clients&.each(&:shutdown)
+    end
+
     private
 
     def scheme
@@ -59,11 +63,13 @@ module Elelem
     end
 
     def mcp_clients
-      config = Pathname.pwd.join(".mcp.json")
-      return [] unless config.exist?
+      @mcp_clients ||= begin
+        config = Pathname.pwd.join(".mcp.json")
+        return [] unless config.exist?
 
-      JSON.parse(config.read).map do |_key, value|
-        MCPClient.new(self, [value["command"]] + value["args"])
+        JSON.parse(config.read).map do |_key, value|
+          MCPClient.new(self, [value["command"]] + value["args"])
+        end
       end
     end
   end
lib/elelem/mcp_client.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "timeout"
+
 module Elelem
   class MCPClient
     attr_reader :tools
@@ -53,6 +55,29 @@ module Elelem
       )
     end
 
+    def shutdown
+      return unless connected?
+
+      configuration.logger.debug("Shutting down MCP client")
+
+      [@stdin, @stdout, @stderr].each do |stream|
+        stream&.close unless stream&.closed?
+      end
+
+      return unless @worker&.alive?
+
+      begin
+        Process.kill("TERM", @worker.pid)
+        # Give it 2 seconds to terminate gracefully
+        Timeout.timeout(2) { @worker.value }
+      rescue Timeout::Error
+        # Force kill if it doesn't respond
+        Process.kill("KILL", @worker.pid) rescue nil
+      rescue Errno::ESRCH
+        # Process already dead
+      end
+    end
+
     private
 
     attr_reader :stdin, :stdout, :stderr, :worker, :configuration