legate 0.1.0

data/lib/legate/mcp/connection/sse.rb

@@ -0,0 +1,292 @@
+# File: lib/legate/mcp/connection/sse.rb
+# frozen_string_literal: true
+
+require 'net/http'
+require 'uri'
+require 'json'
+require_relative '../../errors'
+
+module Legate
+  module Mcp
+    module Connection
+      # Manages a connection to an MCP server via HTTP/SSE.
+      # Uses Server-Sent Events (SSE) for server-to-client notifications
+      # and standard HTTP POST for client-to-server requests/responses.
+      # Automatically reconnects with exponential backoff on stream drops.
+      class Sse
+        MAX_RECONNECT_ATTEMPTS = 5
+        RECONNECT_BASE_DELAY = 1
+        RECONNECT_MAX_DELAY = 30
+
+        attr_reader :url, :last_error, :notification_queue
+
+        def initialize(url:)
+          base_uri = URI.parse(url)
+          base_path = base_uri.path
+          base_path += '/' unless base_path.end_with?('/')
+          @base_uri = base_uri.dup
+          @base_uri.path = base_path
+
+          @sse_uri = @base_uri + 'sse'
+          @message_uri = @base_uri + 'messages'
+          @sse_reader_thread = nil
+          @connected = false
+          @disconnecting = false
+          @request_id_counter = 0
+          @notification_queue = Queue.new
+          @connect_signal = nil
+          @last_error = nil
+          Legate.logger.info("SSE Connection initialized for URL: #{@base_uri}, SSE: #{@sse_uri}, Msg: #{@message_uri}")
+        end
+
+        def connected?
+          @connected && @sse_reader_thread&.alive?
+        end
+
+        # Connects to the SSE endpoint and starts listening for notifications.
+        # Initial connection is synchronous — raises on failure.
+        # After a successful connection, stream drops trigger automatic reconnection.
+        # @raise [ConnectionError] if the initial connection fails.
+        def connect
+          return true if connected?
+
+          Legate.logger.info("Connecting to SSE endpoint: #{@sse_uri}...")
+          @disconnecting = false
+          @last_error = nil
+          @notification_queue.clear
+          @connect_signal = Queue.new
+
+          @sse_reader_thread = Thread.new { connection_loop }
+
+          result = @connect_signal.pop(timeout: 10)
+          @connect_signal = nil
+
+          raise ConnectionError, @last_error || 'Failed to establish SSE connection within timeout' unless result == :connected
+
+          true
+        end
+
+        # Disconnects the SSE stream and stops reconnection attempts.
+        def disconnect
+          return unless @connected || @sse_reader_thread
+
+          Legate.logger.info('Disconnecting SSE connection...')
+          @disconnecting = true
+          @connected = false
+
+          if @sse_reader_thread&.alive?
+            @sse_reader_thread.kill
+            @sse_reader_thread.join(2)
+          end
+
+          @sse_reader_thread = nil
+          @notification_queue.clear
+          Legate.logger.info('SSE connection disconnected.')
+        end
+
+        # Sends a request via HTTP POST and returns the response immediately.
+        # @param json_rpc_hash [Hash] The JSON-RPC request.
+        # @return [Hash] The parsed JSON-RPC response from the server.
+        def send_request(json_rpc_hash)
+          request_json = json_rpc_hash.to_json
+          Legate.logger.debug("-> [MCP Client POST] #{@message_uri} Body: #{request_json}")
+
+          begin
+            http = Net::HTTP.new(@message_uri.hostname, @message_uri.port)
+            http.use_ssl = (@message_uri.scheme == 'https')
+            http.open_timeout = 5
+            http.read_timeout = 15
+
+            request = Net::HTTP::Post.new(@message_uri.request_uri)
+            request['Content-Type'] = 'application/json'
+            request['Accept'] = 'application/json'
+            request.body = request_json
+
+            response = http.request(request)
+
+            unless response.is_a?(Net::HTTPOK)
+              msg = "MCP POST request failed: #{response.code} #{response.message}. Body: #{response.body[0..500]}"
+              Legate.logger.error(msg)
+              @last_error = msg
+              begin
+                error_details = JSON.parse(response.body, symbolize_names: true)
+                if error_details[:error]
+                  raise RemoteToolError.new(error_details[:error][:message], error_details[:error][:code],
+                                            error_details[:error][:data])
+                end
+              rescue JSON::ParserError
+                # Body is not JSON
+              end
+              raise ConnectionError, msg
+            end
+
+            begin
+              response_hash = JSON.parse(response.body, symbolize_names: true)
+              Legate.logger.debug("<- [MCP Client POST Response] #{response_hash.inspect}")
+              response_hash
+            rescue JSON::ParserError => e
+              msg = "Failed to parse MCP JSON response from POST: #{e.message}. Body: #{response.body[0..500]}"
+              Legate.logger.error(msg)
+              @last_error = msg
+              raise ProtocolError, msg
+            end
+          rescue Timeout::Error, Errno::ECONNREFUSED, Errno::EHOSTUNREACH, SocketError => e
+            @last_error = "Failed to send POST to #{@message_uri}: #{e.class} - #{e.message}"
+            Legate.logger.error(@last_error)
+            raise ConnectionError, @last_error
+          rescue Legate::Mcp::ProtocolError, Legate::Mcp::RemoteToolError
+            raise
+          rescue StandardError => e
+            @last_error = "Unexpected error during POST send_request: #{e.class} - #{e.message}"
+            Legate.logger.error("#{@last_error}\n#{e.backtrace.join("\n")}")
+            raise ConnectionError, @last_error
+          end
+        end
+
+        # Reads the next notification from the queue.
+        # @param timeout [Numeric, nil] Seconds to wait, 0 for non-blocking.
+        # @return [Hash, nil] Notification hash or nil if queue is empty/timeout occurs.
+        def read_notification(timeout = 0.1)
+          return nil unless connected?
+
+          begin
+            return @notification_queue.pop(true) if timeout == 0
+
+            @notification_queue.pop(timeout: timeout)
+          rescue ThreadError
+            nil
+          end
+        end
+
+        def next_request_id
+          @request_id_counter += 1
+        end
+
+        private
+
+        def connection_loop
+          attempt = 0
+          initial = true
+
+          until @disconnecting
+            begin
+              attempt += 1
+              run_sse_stream
+              break if @disconnecting
+
+              @connected = false
+              initial = false
+              attempt = 0
+              Legate.logger.info('SSE stream ended, will reconnect.')
+            # IOError covers EOFError, and Timeout::Error covers Net::Open/ReadTimeout,
+            # so the subclasses are omitted here to avoid shadowed (redundant) rescues.
+            rescue ConnectionError, IOError, Errno::ECONNREFUSED,
+                   Errno::ECONNRESET, Errno::EHOSTUNREACH, SocketError,
+                   Timeout::Error => e
+              break if @disconnecting
+
+              @connected = false
+              @last_error = e.message
+
+              if initial
+                signal_connect(:failed)
+                return
+              end
+
+              if attempt > MAX_RECONNECT_ATTEMPTS
+                Legate.logger.error("SSE: Max reconnect attempts (#{MAX_RECONNECT_ATTEMPTS}) reached: #{e.message}")
+                break
+              end
+
+              delay = [RECONNECT_BASE_DELAY * (2**(attempt - 1)), RECONNECT_MAX_DELAY].min
+              Legate.logger.warn("SSE reconnect #{attempt}/#{MAX_RECONNECT_ATTEMPTS} (#{e.class}). Retrying in #{delay}s...")
+              sleep(delay)
+            rescue StandardError => e
+              @connected = false
+              @last_error = "Unexpected SSE error: #{e.class} - #{e.message}"
+              Legate.logger.error("#{@last_error}\n#{e.backtrace&.first(5)&.join("\n")}")
+              signal_connect(:failed) if initial
+              break
+            end
+          end
+
+          @connected = false
+        end
+
+        def run_sse_stream
+          Net::HTTP.start(@sse_uri.hostname, @sse_uri.port,
+                          use_ssl: @sse_uri.scheme == 'https',
+                          open_timeout: 5, read_timeout: 30) do |http|
+            request = Net::HTTP::Get.new(@sse_uri.request_uri)
+            request['Accept'] = 'text/event-stream'
+            request['Cache-Control'] = 'no-cache'
+
+            http.request(request) do |response|
+              unless response.is_a?(Net::HTTPOK) && response['content-type']&.include?('text/event-stream')
+                raise ConnectionError,
+                      "SSE endpoint returned #{response.code}: #{response['content-type']}"
+              end
+
+              Legate.logger.info("SSE connection established with #{@sse_uri}.")
+              @connected = true
+              signal_connect(:connected)
+
+              buffer = +''
+              response.read_body do |chunk|
+                buffer << chunk
+                while (line_end = buffer.index("\n\n"))
+                  event_data = buffer.slice!(0, line_end + 2)
+                  process_sse_event(event_data)
+                end
+              end
+              Legate.logger.info('SSE stream ended.')
+            end
+          end
+        rescue EOFError
+          Legate.logger.info('SSE stream closed by server.')
+        end
+
+        def signal_connect(status)
+          @connect_signal&.push(status)
+          @connect_signal = nil
+        end
+
+        def process_sse_event(event_string)
+          Legate.logger.debug("Processing SSE Event: #{event_string.inspect}")
+          data_buffer = +''
+          event_type = nil
+
+          event_string.each_line do |line|
+            line.chomp!
+            next if line.empty?
+            next if line.start_with?(':')
+
+            field, value = line.split(':', 2)
+            value&.strip!
+
+            case field
+            when 'event'
+              event_type = value
+            when 'data'
+              data_buffer << value << "\n"
+            when 'retry'
+              Legate.logger.debug("Received SSE retry suggestion: #{value}ms")
+            else
+              Legate.logger.warn("Ignoring unknown SSE field: #{field}")
+            end
+          end
+
+          return if data_buffer.empty?
+
+          begin
+            message = JSON.parse(data_buffer, symbolize_names: true)
+            Legate.logger.debug("Parsed SSE notification (event: #{event_type || 'message'}): #{message.inspect}")
+            @notification_queue << message
+          rescue JSON::ParserError => e
+            Legate.logger.error("Failed to parse JSON from SSE data: #{e.message}. Data: #{data_buffer.inspect}")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legate/mcp/connection/stdio.rb

@@ -0,0 +1,273 @@
+# File: lib/legate/mcp/connection/stdio.rb
+# frozen_string_literal: true
+
+require 'open3'
+require 'json'
+require_relative '../../errors'
+
+module Legate
+  module Mcp
+    module Connection
+      # Manages a connection to an MCP server via STDIO.
+      class Stdio
+        # How long to wait for process startup or initial output
+        PROCESS_START_TIMEOUT = 5 # seconds
+        # How long to wait when reading a response line
+        READ_TIMEOUT = 10 # seconds
+        # Max consecutive JSON parse errors before considering connection broken
+        PARSE_ERROR_THRESHOLD = 5
+
+        attr_reader :command, :args, :last_error
+
+        def initialize(command:, args: [])
+          @command = command
+          @args = args
+          @stdin = nil
+          @stdout = nil
+          @stderr = nil
+          @wait_thr = nil
+          @stderr_thread = nil
+          @connected = false
+          @request_id_counter = 0
+          @response_queue = Queue.new # <-- Back to response_queue
+          @notification_queue = Queue.new # <-- Back to notification_queue
+          @stdout_reader_thread = nil
+          @last_error = nil
+          @pid = nil
+          @consecutive_parse_errors = 0 # Initialize counter
+        end
+
+        def connected?
+          @connected && @wait_thr&.alive?
+        end
+
+        # Connects to the MCP server process.
+        # Launches the command and starts threads to monitor stdout/stderr.
+        # @raise [ConnectionError] if the process fails to start or terminates unexpectedly.
+        def connect
+          return true if connected?
+
+          Mcp.logger.info("Connecting via STDIO: #{@command} #{@args.join(' ')}")
+          @last_error = nil
+          stderr_pipe_read, stderr_pipe_write = IO.pipe
+
+          begin
+            # Use popen3 to capture stdin, stdout, stderr, and wait_thr
+            @stdin, @stdout, @stderr, @wait_thr = Open3.popen3(@command, *@args, err: stderr_pipe_write)
+            @pid = @wait_thr.pid
+            stderr_pipe_write.close # Close the write end in the parent
+
+            Mcp.logger.debug("MCP process started with PID: #{@pid}")
+
+            # Thread to read stderr and log/store errors
+            @stderr_thread = Thread.new do
+              stderr_pipe_read.each_line do |line|
+                Mcp.logger.error("[MCP Server STDERR] #{line.chomp}")
+                @last_error = line.chomp # Store last error line
+              end
+            rescue IOError => e
+              Mcp.logger.debug("Stderr pipe closed: #{e.message}")
+            ensure
+              stderr_pipe_read.close unless stderr_pipe_read.closed?
+            end
+
+            # Thread to continuously read stdout and parse JSON-RPC messages
+            @stdout_reader_thread = Thread.new do
+              Mcp.logger.debug("[Stdio Connection #{@pid}] stdout_reader_thread starting...")
+              begin
+                @stdout.each_line do |line|
+                  # Handle potential encoding issues from subprocess output
+                  line.force_encoding('UTF-8')
+                  line.scrub!('') # Remove invalid bytes
+                  line.strip! # Remove leading/trailing whitespace
+                  next if line.empty? # Skip empty lines
+
+                  Legate.logger.debug("<- [MCP Server STDOUT Raw] #{line}")
+
+                  # Attempt to parse only if it looks like JSON
+                  if line.start_with?('{') || line.start_with?('[')
+                    begin
+                      message = JSON.parse(line, symbolize_names: true)
+                      Legate.logger.debug("[Stdio Connection #{@pid}] Received Parsed JSON:\n#{JSON.pretty_generate(message)}")
+                      @consecutive_parse_errors = 0 # Reset on successful parse
+
+                      # Route to correct queue based on ID presence/value
+                      if message.key?(:id) && message[:id].nil? # MCP Notifications might have null id
+                        @notification_queue << message
+                      elsif message.key?(:id)
+                        Legate.logger.debug("[Stdio Connection #{@pid}] Queuing response ID: #{message[:id]}")
+                        @response_queue << message # Responses have non-null id
+                      else # Assume notification for now
+                        @notification_queue << message
+                        Legate.logger.warn("Received MCP message without explicit id: #{message.inspect}")
+                      end
+                    rescue JSON::ParserError => e
+                      Legate.logger.error("Failed to parse potential MCP JSON from stdout: #{e.message}. Line: #{line}")
+                      @consecutive_parse_errors += 1
+                      if @consecutive_parse_errors >= PARSE_ERROR_THRESHOLD # Use >= for clarity
+                        Legate.logger.fatal("Too many consecutive JSON parse errors (#{PARSE_ERROR_THRESHOLD} reached). Assuming MCP connection broken.")
+                        @connected = false # Mark connection as broken
+                        @last_error = 'Too many consecutive JSON parse errors.'
+                        break # Stop reading from stdout
+                      end
+                    end
+                  else
+                    # Log lines that don't look like JSON instead of trying to parse
+                    Legate.logger.debug("Skipping non-JSON line from MCP STDOUT: #{line}")
+                    # Do not increment parse error count for these lines
+                  end
+                end
+                Legate.logger.info('MCP Server stdout stream ended.')
+              rescue IOError => e
+                Legate.logger.info("MCP Server stdout pipe closed: #{e.message}")
+              ensure
+                @connected = false # Mark as disconnected if stdout closes or loop breaks due to errors
+                Mcp.logger.debug("[Stdio Connection #{@pid}] stdout_reader_thread finished.")
+              end
+            end
+
+            @connected = true
+            Mcp.logger.info('MCP STDIO connection established.')
+            true
+          rescue Errno::ENOENT => e
+            @last_error = "Command not found: #{@command}"
+            Mcp.logger.error("#{@last_error} - #{e.message}")
+            raise ConnectionError, @last_error
+          rescue StandardError => e
+            @last_error = "Failed to start MCP process: #{e.message}"
+            Mcp.logger.error("#{@last_error}")
+            # Clean up if process started partially
+            disconnect(force: true)
+            raise ConnectionError, @last_error
+          end
+        end
+
+        # Sends a JSON-RPC request object to the server process.
+        # @param json_rpc_hash [Hash] The request hash (e.g., {jsonrpc: '2.0', method: '...', params: ..., id: ...})
+        # @raise [ConnectionError] if not connected.
+        def send_request(json_rpc_hash)
+          raise ConnectionError, 'Not connected' unless connected?
+
+          begin
+            request_json = json_rpc_hash.to_json
+            Mcp.logger.debug("-> [MCP Client STDIN] #{request_json}")
+            @stdin.puts(request_json)
+            @stdin.flush # Ensure data is sent immediately
+          rescue Errno::EPIPE => e
+            @connected = false
+            @last_error = "MCP process stdin pipe broke: #{e.message}"
+            Mcp.logger.error(@last_error)
+            raise ConnectionError, @last_error
+          rescue StandardError => e
+            @connected = false
+            @last_error = "Error writing to MCP process stdin: #{e.class} - #{e.message}"
+            Mcp.logger.error("#{@last_error}\n#{e.backtrace.join("\n")}")
+            raise ConnectionError, @last_error
+          end
+        end
+
+        # Reads the next available response or notification.
+        # This is a low-level method; typically use Client methods which match request/response.
+        # @param timeout [Numeric, nil] Seconds to wait for a message, nil to wait indefinitely.
+        # @return [Hash, nil] The parsed JSON-RPC message, or nil if timeout occurs.
+        # @raise [ConnectionError] if not connected or connection lost.
+        def read_message(timeout = READ_TIMEOUT)
+          raise ConnectionError, 'Not connected' unless connected?
+
+          # Check both queues, prioritize responses if available
+          begin
+            @response_queue.pop(true) # non_block = true
+          rescue ThreadError
+            # Response queue empty, check notifications
+            begin
+              @notification_queue.pop(true)
+            rescue ThreadError
+              # Both empty, wait with timeout if specified
+              return nil if timeout == 0 # Don't wait if timeout is 0
+
+              deadline = timeout ? Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout : nil
+              loop do
+                remaining = deadline ? deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC) : 1.0
+                return nil if remaining <= 0
+
+                wait = [remaining, 0.5].min
+                msg = @response_queue.pop(timeout: wait)
+                return msg if msg
+
+                begin
+                  return @notification_queue.pop(true)
+                rescue ThreadError
+                  # notification queue empty too
+                end
+
+                raise ConnectionError, 'Connection lost while waiting for message' unless connected?
+              end
+            end
+          end
+        end
+
+        # Disconnects from the server process.
+        # Terminates the process and cleans up threads.
+        # @param force [Boolean] If true, use SIGKILL if SIGTERM fails.
+        # @param timeout [Numeric] Seconds to wait for graceful shutdown.
+        def disconnect(force: false, timeout: 5)
+          return unless @connected || @wait_thr # Only proceed if we have something to disconnect
+
+          Mcp.logger.info("Disconnecting from MCP STDIO process (PID: #{@pid})...")
+          @connected = false
+
+          # Close stdin to signal EOF to the process
+          @stdin&.close unless @stdin&.closed?
+
+          # Close stdout/stderr BEFORE killing reader threads so they
+          # unblock from IO.read and can exit cleanly
+          @stdout&.close unless @stdout&.closed?
+          @stderr&.close unless @stderr&.closed?
+
+          # Join reader threads with timeout, then force-kill if stuck
+          [@stdout_reader_thread, @stderr_thread].each do |thr|
+            next unless thr&.alive?
+
+            thr.join(2) || thr.kill
+          end
+
+          # Terminate the child process
+          if @wait_thr&.pid
+            pid = @wait_thr.pid
+            begin
+              Mcp.logger.debug("Sending SIGTERM to PID #{pid}...")
+              Process.kill('TERM', pid)
+              process_exited = @wait_thr.join(timeout)
+
+              unless process_exited
+                Mcp.logger.warn("MCP process PID #{pid} did not exit after SIGTERM and #{timeout}s timeout.")
+                if force
+                  Mcp.logger.warn("Forcing shutdown with SIGKILL for PID #{pid}.")
+                  Process.kill('KILL', pid)
+                  @wait_thr.join(2)
+                end
+              end
+              Mcp.logger.info("MCP process PID #{pid} terminated. Status: #{@wait_thr.value}")
+            rescue Errno::ESRCH
+              Mcp.logger.info("MCP process PID #{pid} already exited.")
+            rescue StandardError => e
+              Mcp.logger.debug("Caught StandardError during termination: #{e.class}")
+              Mcp.logger.error("Error during process termination for PID #{pid}: #{e.message}")
+            end
+          end
+
+          @stdin = @stdout = @stderr = @wait_thr = @pid = nil
+          @response_queue.clear
+          @notification_queue.clear
+          Mcp.logger.info('MCP STDIO connection closed.')
+        end
+
+        # Generates the next unique request ID.
+        # @return [Integer]
+        def next_request_id
+          @request_id_counter += 1
+        end
+      end
+    end
+  end
+end

data/lib/legate/mcp/connection_manager.rb

@@ -0,0 +1,103 @@
+# File: lib/legate/mcp/connection_manager.rb
+# frozen_string_literal: true
+
+require_relative 'client'
+require_relative 'tool_wrapper'
+
+module Legate
+  module Mcp
+    # Owns an agent's MCP client connections: connecting to configured servers,
+    # discovering and registering their tools into the agent's tool registry, and
+    # disconnecting. Extracted from Legate::Agent to keep MCP lifecycle out of the
+    # agent's core responsibilities.
+    class ConnectionManager
+      attr_reader :clients
+
+      # @param tool_registry [Legate::ToolRegistry] the agent's registry that MCP tools register into
+      # @param selected_tool_names [Array<Symbol>] tool names the agent selected (others are skipped)
+      # @param agent_name [Symbol, String] for log context
+      def initialize(tool_registry:, selected_tool_names:, agent_name:)
+        @tool_registry = tool_registry
+        @selected_tool_names = selected_tool_names
+        @agent_name = agent_name
+        @clients = []
+      end
+
+      # Connects to each configured MCP server and registers its selected tools.
+      # @param servers_config [Array<Hash>, nil] server configs from the definition
+      def connect(servers_config)
+        return if servers_config.nil? || servers_config.empty?
+
+        servers_config.each do |config|
+          # Transform keys to symbols for the client
+          symbolized_config = config.transform_keys(&:to_sym)
+          Legate.logger.info("Attempting to connect to MCP server: #{symbolized_config.inspect}")
+          begin
+            unless %w[stdio sse].include?(symbolized_config[:type])
+              Legate.logger.error("Unsupported MCP server type specified: #{symbolized_config[:type].inspect}. Skipping configuration: #{symbolized_config.inspect}")
+              next # Skip to the next server config
+            end
+
+            # Explicitly convert known string type values to symbols
+            if symbolized_config[:type] == 'stdio'
+              symbolized_config[:type] = :stdio
+            elsif symbolized_config[:type] == 'sse'
+              symbolized_config[:type] = :sse
+            end
+            # Pass the modified hash
+            client = Legate::Mcp::Client.new(symbolized_config)
+            client.connect # This performs handshake and gets capabilities
+            @clients << client
+            discover_and_register_tools(client)
+          rescue Legate::Mcp::ConnectionError, Legate::Mcp::ProtocolError => e # More specific MCP errors
+            Legate.logger.error("Failed to connect or handshake with MCP server #{config.inspect}: #{e.message}")
+          rescue Legate::Mcp::Error => e
+            Legate.logger.error("MCP-related error connecting to server #{config.inspect}: #{e.message}")
+          rescue StandardError => e
+            Legate.logger.error("Unexpected error connecting to MCP server #{config.inspect}: #{e.class} - #{e.message}")
+          end
+        end
+      end
+
+      # Disconnects all active MCP clients.
+      def disconnect
+        return if @clients.nil? || @clients.empty?
+
+        @clients.each do |client|
+          Legate.logger.info('Disconnecting MCP client...')
+          client.disconnect
+        rescue StandardError => e
+          Legate.logger.error("Error disconnecting MCP client: #{e.message}")
+        end
+        @clients.clear
+      end
+
+      private
+
+      # Discovers tools from a connected MCP client and registers the selected
+      # ones with the agent's registry.
+      # @param client [Legate::Mcp::Client]
+      def discover_and_register_tools(client)
+        Legate.logger.debug("[Agent E2E Debug] discover_and_register - @tool_registry ID: #{@tool_registry.object_id}")
+        begin
+          mcp_tool_schemas = client.list_tools
+          Legate.logger.debug("[Agent E2E Debug] list_tools returned: #{mcp_tool_schemas.inspect}")
+          Legate.logger.info("Discovered #{mcp_tool_schemas.count} tools from MCP server.")
+          mcp_tool_schemas.each do |schema|
+            tool_name_sym = schema[:name].to_sym
+            if @selected_tool_names.include?(tool_name_sym)
+              # Pass the agent's specific registry instance
+              Legate::Mcp::ToolWrapper.from_mcp_schema(schema, client, @tool_registry)
+            else
+              Legate.logger.debug("Skipping registration of MCP tool '#{tool_name_sym}' as it was not selected in agent definition.")
+            end
+          end
+        rescue Legate::Mcp::Error => e
+          Legate.logger.error("Failed to list tools from MCP server: #{e.message}")
+        rescue StandardError => e
+          Legate.logger.error("Unexpected error discovering MCP tools: #{e.class} - #{e.message}")
+        end
+      end
+    end
+  end
+end