ruby_llm-mcp 0.7.1 → 0.8.0

data/lib/generators/ruby_llm/mcp/oauth/templates/controllers/mcp_connections_controller.rb.tt

@@ -0,0 +1,239 @@
+# frozen_string_literal: true
+
+# Controller for managing MCP OAuth connections
+<% if namespace_name -%>
+class <%= namespace_name %>::McpConnectionsController < ApplicationController
+<% else -%>
+class McpConnectionsController < ApplicationController
+<% end -%>
+  include ActionView::Helpers::DateHelper
+
+  before_action :<%= authenticate_method %>
+
+  OAUTH_FLOW_TIMEOUT_SECONDS = 10.minutes.to_i
+
+  def index
+    @credentials = <%= current_user_method %>.mcp_oauth_credentials.order(created_at: :desc)
+  end
+
+  def show
+    @credential = <%= current_user_method %>.mcp_oauth_credentials.find(params[:id])
+
+    begin
+      client = McpClient.for(<%= current_user_method %>, server_url: @credential.server_url)
+      @tools = client.tools
+
+      Rails.logger.info "Loaded #{@tools.count} tools from MCP server: #{@credential.server_url}"
+    rescue McpClient::NotAuthenticatedError => e
+      Rails.logger.error "MCP authentication error: #{e.message}"
+      redirect_to mcp_connections_path,
+                  alert: "Authentication failed. Please refresh your connection."
+      nil
+    rescue StandardError => e
+      Rails.logger.error "Failed to retrieve MCP tools: #{e.message}"
+      redirect_to mcp_connections_path,
+                  alert: "Failed to connect to MCP server: #{e.message}"
+      nil
+    ensure
+      client&.stop
+    end
+  rescue ActiveRecord::RecordNotFound
+    redirect_to mcp_connections_path,
+                alert: "Connection not found"
+  end
+
+  def create
+    name = params[:name]
+    server_url = params[:server_url]
+
+    unless name.present?
+      redirect_to mcp_connections_path, alert: "Please provide a server name"
+      return
+    end
+
+    unless server_url.present?
+      redirect_to mcp_connections_path, alert: "Please provide an MCP server URL"
+      return
+    end
+
+    scope = params[:scope] || "mcp:read mcp:write"
+
+    # Create OAuth provider with user's token storage
+    # Storage will automatically save name and scope to McpOauthState table
+    oauth_provider = RubyLLM::MCP::Auth::OAuthProvider.new(
+      server_url: server_url,
+      redirect_uri: mcp_connections_callback_url,
+      scope: scope,
+      storage: token_storage(server_url, name, scope),
+      logger: Rails.logger
+    )
+
+    # Start OAuth flow
+    begin
+      auth_url = oauth_provider.start_authorization_flow
+      # OAuth state (including name and scope) is now stored in database via token_storage
+      # No need to use session - database persists across redirects
+
+      redirect_to auth_url, allow_other_host: true
+    rescue StandardError => e
+      Rails.logger.error "MCP OAuth flow start failed: #{e.message}"
+      redirect_to mcp_connections_path,
+                  alert: "Failed to start OAuth flow: #{e.message}"
+    end
+  end
+
+  def update
+    credential = <%= current_user_method %>.mcp_oauth_credentials.find(params[:id])
+
+    oauth_provider = RubyLLM::MCP::Auth::OAuthProvider.new(
+      server_url: credential.server_url,
+      storage: token_storage(credential.server_url, credential.name, credential.token&.scope),
+      logger: Rails.logger
+    )
+
+    refreshed_token = oauth_provider.access_token
+
+    if refreshed_token
+      credential.update!(last_refreshed_at: Time.current)
+
+      Rails.logger.info "Token refreshed for <%= user_variable_name %> #{<%= current_user_method %>.id}, server: #{credential.server_url}"
+
+      redirect_to mcp_connections_path,
+                  notice: "Token refreshed successfully. Expires #{time_ago_in_words(credential.token_expires_at)} from now."
+    else
+      Rails.logger.warn "Token refresh failed for <%= user_variable_name %> #{<%= current_user_method %>.id}, server: #{credential.server_url}"
+
+      redirect_to mcp_connections_path,
+                  alert: "Token refresh failed. Please reconnect to authorize again."
+    end
+  rescue ActiveRecord::RecordNotFound
+    redirect_to mcp_connections_path,
+                alert: "Connection not found"
+  rescue StandardError => e
+    Rails.logger.error "Token refresh error: #{e.message}"
+    redirect_to mcp_connections_path,
+                alert: "Token refresh failed: #{e.message}"
+  end
+
+  def destroy
+    credential = <%= current_user_method %>.mcp_oauth_credentials.find(params[:id])
+    server_url = credential.server_url
+    credential.destroy
+
+    Rails.logger.info "<%= user_model_name %> #{<%= current_user_method %>.id} disconnected from MCP server: #{server_url}"
+
+    redirect_to mcp_connections_path,
+                notice: "MCP server disconnected successfully"
+  rescue ActiveRecord::RecordNotFound
+    redirect_to mcp_connections_path,
+                alert: "Connection not found"
+  end
+
+  def callback
+    return if oauth_error_present?
+
+    oauth_state = retrieve_and_validate_oauth_state
+    return unless oauth_state
+
+    complete_oauth_flow_for_user(oauth_state)
+  end
+
+  private
+
+  def retrieve_and_validate_oauth_state
+    # Look up OAuth state from database using the state parameter from callback
+    state_param = params[:state]
+
+    unless state_param
+      redirect_to mcp_connections_path, alert: "Missing OAuth state parameter"
+      return nil
+    end
+
+    oauth_state = <%= current_user_method %>.mcp_oauth_states.find_by(state_param: state_param)
+
+    unless oauth_state
+      redirect_to mcp_connections_path, alert: "OAuth state not found. Please try again."
+      return nil
+    end
+
+    if oauth_state.expires_at < Time.current
+      oauth_state.destroy
+      redirect_to mcp_connections_path, alert: "OAuth flow timed out. Please try again."
+      return nil
+    end
+
+    oauth_state
+  end
+
+  def oauth_error_present?
+    return false unless params[:error]
+
+    error_message = params[:error_description] || params[:error]
+    redirect_to mcp_connections_path, alert: "OAuth authorization failed: #{error_message}"
+    true
+  end
+
+  def complete_oauth_flow_for_user(oauth_state)
+    oauth_provider = create_oauth_provider_from_state(oauth_state)
+
+    oauth_provider.complete_authorization_flow(params[:code], params[:state])
+
+    # Clean up the OAuth state record after successful completion
+    oauth_state.destroy
+
+    log_successful_oauth(oauth_state)
+    redirect_after_success
+  rescue StandardError => e
+    handle_oauth_callback_error(e, oauth_state)
+  end
+
+  def create_oauth_provider_from_state(oauth_state)
+    RubyLLM::MCP::Auth::OAuthProvider.new(
+      server_url: oauth_state.server_url,
+      redirect_uri: mcp_connections_callback_url,
+      scope: oauth_state.scope,
+      storage: token_storage(oauth_state.server_url, oauth_state.name, oauth_state.scope),
+      logger: Rails.logger
+    )
+  end
+
+  def log_successful_oauth(oauth_state)
+    Rails.logger.info "MCP OAuth completed for <%= user_variable_name %> #{<%= current_user_method %>.id}, " \
+                      "server: #{oauth_state.server_url}"
+  end
+
+  def redirect_after_success
+    redirect_to mcp_connections_path, notice: "Successfully connected to MCP server!"
+  end
+
+  def handle_oauth_callback_error(error, oauth_state = nil)
+    Rails.logger.error "MCP OAuth callback failed: #{error.message}"
+    oauth_state&.destroy # Clean up failed state
+    redirect_to mcp_connections_path, alert: "OAuth authorization failed: #{error.message}"
+  end
+
+  def token_storage(server_url, name = nil, scope = nil)
+    <%= current_user_method %>.mcp_token_storage(server_url, name, scope)
+  end
+
+  def mcp_connections_callback_url
+<% if namespace_name -%>
+    callback_<%= namespace_name.underscore %>_mcp_connections_url
+<% else -%>
+    callback_mcp_connections_url
+<% end -%>
+  end
+<% if namespace_name -%>
+
+  # Route helper methods for namespaced routes
+  def mcp_connections_path
+    <%= namespace_name.underscore %>_mcp_connections_path
+  end
+
+  def mcp_connection_path(credential)
+    <%= namespace_name.underscore %>_mcp_connection_path(credential)
+  end
+
+  helper_method :mcp_connections_path, :mcp_connection_path
+<% end -%>
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/jobs/cleanup_expired_oauth_states_job.rb.tt

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+# Background job to cleanup expired OAuth states
+# Schedule this to run hourly or daily to prevent state table bloat
+#
+# Example with sidekiq-scheduler:
+#   cleanup_expired_states:
+#     cron: '0 * * * *'  # Every hour
+#     class: CleanupExpiredOauthStatesJob
+#
+# Example with whenever gem:
+#   every 1.hour do
+#     runner "CleanupExpiredOauthStatesJob.perform_later"
+#   end
+class CleanupExpiredOauthStatesJob < ApplicationJob
+  queue_as :default
+
+  # Cleanup expired OAuth flow states
+  # States are temporary and only needed during the OAuth flow (typically < 10 minutes)
+  def perform
+    deleted_count = <%= state_model_name %>.cleanup_expired
+
+    Rails.logger.info "Cleaned up #{deleted_count} expired MCP OAuth states"
+
+    deleted_count
+  end
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/jobs/example_job.rb.tt

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+# Example background job using MCP with per-<%= user_variable_name %> OAuth
+class AiResearchJob < ApplicationJob
+  queue_as :default
+
+  # Retry on authentication errors (<%= user_variable_name %> might reconnect)
+  retry_on McpClient::NotAuthenticatedError,
+           wait: 1.hour,
+           attempts: 3 do |job, _exception|
+    # After retries exhausted, notify <%= user_variable_name %>
+    <%= user_variable_name %> = <%= user_model_name %>.find(job.arguments.first)
+    <%= user_model_name %>Mailer.mcp_auth_required(<%= user_variable_name %>).deliver_now
+  end
+
+  # Retry on transient network errors
+  retry_on RubyLLM::MCP::Errors::TransportError,
+           wait: :exponentially_longer,
+           attempts: 5
+
+  # Perform AI research using <%= user_variable_name %>'s MCP connection
+  # @param <%= user_variable_name %>_id [Integer] ID of the <%= user_variable_name %>
+  # @param server_url [String] MCP server URL to use
+  # @param query [String] research query
+  # @param options [Hash] additional options
+  def perform(<%= user_variable_name %>_id, server_url, query, options = {})
+    <%= user_variable_name %> = <%= user_model_name %>.find(<%= user_variable_name %>_id)
+
+    # Create MCP client with <%= user_variable_name %>'s OAuth token
+    client = McpClient.for(<%= user_variable_name %>, server_url: server_url)
+
+    begin
+      # Get tools with <%= user_variable_name %>'s permissions
+      tools = client.tools
+      Rails.logger.info "Loaded #{tools.count} MCP tools for <%= user_variable_name %> #{<%= user_variable_name %>_id}"
+
+      # Create AI chat with <%= user_variable_name %>'s MCP context
+      chat = RubyLLM.chat(provider: options[:provider] || "anthropic/claude-sonnet-4")
+                    .with_tools(*tools)
+
+      # Execute research
+      response = chat.ask(query)
+
+      # Save results
+      save_research_results(<%= user_variable_name %>, query, response)
+
+      # Notify <%= user_variable_name %> of completion
+      notify_completion(<%= user_variable_name %>, query)
+
+      Rails.logger.info "Completed AI research for <%= user_variable_name %> #{<%= user_variable_name %>_id}"
+    ensure
+      # Always close client connection
+      client&.stop
+    end
+  end
+
+  private
+
+  def save_research_results(<%= user_variable_name %>, query, response)
+    # Customize based on your schema
+    <%= user_variable_name %>.research_results.create!(
+      query: query,
+      result: response.text,
+      completed_at: Time.current
+    )
+  end
+
+  def notify_completion(<%= user_variable_name %>, query)
+    # Notify via email, ActionCable, or other mechanism
+    ActionCable.server.broadcast(
+      "<%= user_variable_name %>_#{<%= user_variable_name %>.id}_notifications",
+      {
+        type: "research_complete",
+        message: "Research completed: #{query.truncate(50)}"
+      }
+    )
+  end
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/lib/mcp_client.rb.tt

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+# Creates per-<%= user_variable_name %> MCP clients with OAuth authentication
+class McpClient
+  class NotAuthenticatedError < StandardError; end
+
+  # Create MCP client for a specific <%= user_variable_name %>
+  # @param <%= user_variable_name %> [<%= user_model_name %>] the <%= user_variable_name %> to create client for
+  # @param server_url [String] MCP server URL (required)
+  # @param scope [String] OAuth scopes to request
+  # @return [RubyLLM::MCP::Client] configured MCP client
+  # @raise [NotAuthenticatedError] if <%= user_variable_name %> hasn't connected to MCP server
+  # @raise [ArgumentError] if server_url not provided
+  def self.for(<%= user_variable_name %>, server_url:, scope: nil)
+    raise ArgumentError, "server_url is required" unless server_url.present?
+    scope ||= "mcp:read mcp:write"
+
+    unless <%= user_variable_name %>.mcp_connected?(server_url)
+      raise NotAuthenticatedError,
+            "<%= user_model_name %> #{<%= user_variable_name %>.id} has not connected to MCP server: #{server_url}. " \
+            "Please complete OAuth flow first."
+    end
+
+    storage = <%= user_variable_name %>.mcp_token_storage(server_url)
+
+    RubyLLM::MCP.client(
+      name: "<%= user_variable_name %>-#{<%= user_variable_name %>.id}-#{server_url.hash.abs}",
+      transport_type: determine_transport_type(server_url),
+      config: {
+        url: server_url,
+        oauth: {
+          storage: storage,
+          scope: scope
+        }
+      }
+    )
+  end
+
+  # Create MCP client for <%= user_variable_name %>, returning nil if not authenticated
+  # @param <%= user_variable_name %> [<%= user_model_name %>] the <%= user_variable_name %>
+  # @param server_url [String] MCP server URL (required)
+  # @return [RubyLLM::MCP::Client, nil] client or nil
+  def self.for_with_fallback(<%= user_variable_name %>, server_url:)
+    self.for(<%= user_variable_name %>, server_url: server_url)
+  rescue NotAuthenticatedError, ArgumentError
+    nil
+  end
+
+  # Check if <%= user_variable_name %> has valid MCP connection
+  # @param <%= user_variable_name %> [<%= user_model_name %>] the <%= user_variable_name %>
+  # @param server_url [String] MCP server URL (required)
+  # @return [Boolean] true if <%= user_variable_name %> has valid token
+  def self.connected?(<%= user_variable_name %>, server_url:)
+    return false unless server_url.present?
+
+    credential = <%= user_variable_name %>.mcp_oauth_credentials.find_by(server_url: server_url)
+    credential&.valid_token? || false
+  end
+
+  # Determine transport type from URL
+  # @param url [String] server URL
+  # @return [Symbol] :sse or :streamable
+  def self.determine_transport_type(url)
+    url.include?("/sse") ? :sse : :streamable
+  end
+
+  private_class_method :determine_transport_type
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/migrations/create_mcp_oauth_credentials.rb.tt

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+class CreateMcpOauthCredentials < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :<%= credential_table_name %> do |t|
+      t.references :<%= user_table_name.singularize %>, null: false, foreign_key: true, index: true
+      t.string :name, null: false
+      t.string :server_url, null: false
+      t.text :token_data
+      t.text :client_info_data
+      t.datetime :token_expires_at
+      t.datetime :last_refreshed_at
+
+      t.timestamps
+
+      t.index [:<%= user_table_name.singularize %>_id, :server_url], unique: true, name: "index_mcp_oauth_on_<%= user_table_name.singularize %>_and_server"
+    end
+  end
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/migrations/create_mcp_oauth_states.rb.tt

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+class CreateMcpOauthStates < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :<%= state_table_name %> do |t|
+      t.references :<%= user_table_name.singularize %>, null: false, foreign_key: true
+      t.string :server_url, null: false
+      t.string :state_param, null: false
+      t.text :pkce_data, null: false
+      t.datetime :expires_at, null: false
+      t.string :name
+      t.string :scope
+
+      t.timestamps
+
+      t.index [:<%= user_table_name.singularize %>_id, :state_param], unique: true, name: "index_mcp_oauth_states_on_<%= user_table_name.singularize %>_and_state"
+      t.index [:<%= user_table_name.singularize %>_id, :server_url], unique: true, name: "index_mcp_oauth_states_on_<%= user_table_name.singularize %>_and_server"
+      t.index :expires_at, name: "index_mcp_oauth_states_on_expires_at"
+    end
+  end
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/models/mcp_oauth_credential.rb.tt

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+class <%= credential_model_name %> < ApplicationRecord
+  belongs_to :<%= user_table_name.singularize %>
+
+  encrypts :token_data
+  encrypts :client_info_data
+
+  validates :name, presence: true,
+                   uniqueness: { scope: :<%= user_table_name.singularize %>_id }
+  validates :server_url, presence: true,
+                         uniqueness: { scope: :<%= user_table_name.singularize %>_id },
+                         format: { with: URI::DEFAULT_PARSER.make_regexp(%w[http https]) }
+
+  # Get token object from stored data
+  def token
+    return nil unless token_data.present?
+
+    RubyLLM::MCP::Auth::Token.from_h(JSON.parse(token_data, symbolize_names: true))
+  end
+
+  # Set token and update expiration
+  def token=(token_obj)
+    self.token_data = token_obj.to_h.to_json
+    self.token_expires_at = token_obj.expires_at
+  end
+
+  # Get client info object from stored data
+  def client_info
+    return nil unless client_info_data.present?
+
+    RubyLLM::MCP::Auth::ClientInfo.from_h(JSON.parse(client_info_data, symbolize_names: true))
+  end
+
+  # Set client info
+  def client_info=(info_obj)
+    self.client_info_data = info_obj.to_h.to_json
+  end
+
+  # Check if token is expired
+  def expired?
+    token&.expired?
+  end
+
+  # Check if token expires soon (within 5 minutes)
+  def expires_soon?
+    token&.expires_soon?
+  end
+
+  # Check if token is valid
+  def valid_token?
+    token && !token.expired?
+  end
+end

data/lib/generators/ruby_llm/mcp/oauth/templates/models/mcp_oauth_state.rb.tt

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+class <%= state_model_name %> < ApplicationRecord
+  belongs_to :<%= user_table_name.singularize %>
+
+  encrypts :pkce_data
+
+  validates :state_param, presence: true
+  validates :expires_at, presence: true
+
+  scope :expired, -> { where("expires_at < ?", Time.current) }
+  scope :for_<%= user_table_name.singularize %>, ->(<%= user_table_name.singularize %>_id) { where(<%= user_table_name.singularize %>_id: <%= user_table_name.singularize %>_id) }
+
+  # Clean up expired OAuth flow states
+  def self.cleanup_expired
+    expired.delete_all
+  end
+
+  # Get PKCE object from stored data
+  def pkce
+    return nil unless pkce_data.present?
+
+    RubyLLM::MCP::Auth::PKCE.from_h(JSON.parse(pkce_data, symbolize_names: true))
+  end
+
+  # Set PKCE object
+  def pkce=(pkce_obj)
+    self.pkce_data = pkce_obj.to_h.to_json
+  end
+end