mcp_toolkit 0.3.0 → 0.4.0

data/lib/mcp_toolkit/dispatcher.rb

@@ -0,0 +1,234 @@
+# frozen_string_literal: true
+
+# Hand-rolled JSON-RPC dispatcher for the AUTHORITY + gateway path: it handles a
+# single JSON-RPC request (one element of a batch), serving the host's own tools
+# and, as a gateway, aggregating + proxying upstream MCP servers — WITHOUT the
+# official `mcp` SDK in the request path.
+#
+# The gem carries two dispatch front-ends by design (see McpToolkit::Protocol):
+# this Dispatcher for the authority/gateway endpoint, and McpToolkit::Server.build
+# (the SDK wrapper) for satellites. They are independent; nothing here touches the
+# satellite path.
+#
+# Everything host-specific is injected:
+#   * `context`  — an McpToolkit::Authority::Context (the resolved account, the
+#                  authenticated principal, and the bearer token to forward
+#                  upstream). Re-created PER JSON-RPC request by the transport, so
+#                  each batch element carries its own account.
+#   * `config`   — server identity (`server_name`/`server_version`), the
+#                  negotiable protocol versions, the registered `upstreams`, and
+#                  the `tool_provider` (the host's api-agnostic tool catalog).
+#
+# The wire behavior — top-level JSON-RPC tool-error codes, `initialize`
+# capabilities `{ tools: { listChanged: true } }`, 3-version negotiation, verbatim
+# upstream error relay, and the custom `notifications/<app>/tools/list_changed`
+# cache-bust — is the byte contract of a first-party endpoint and is preserved
+# exactly.
+class McpToolkit::Dispatcher
+  attr_reader :context, :config
+
+  def initialize(context:, config: McpToolkit.config)
+    @context = context
+    @config = config
+  end
+
+  def handle_request(request)
+    dispatch_request(request)
+  rescue McpToolkit::Protocol::Error => e
+    return nil unless request.key?("id")
+
+    McpToolkit::Protocol.error_response(id: request["id"], error: e)
+  rescue StandardError => e
+    config.logger&.error("MCP dispatcher error: #{e.message}\n#{e.backtrace&.join("\n")}")
+    return nil unless request.key?("id")
+
+    McpToolkit::Protocol.error_response(
+      id: request["id"],
+      error: McpToolkit::Protocol::InternalError.new(e.message)
+    )
+  end
+
+  private
+
+  # Happy path; raises here are turned into JSON-RPC errors by handle_request.
+  def dispatch_request(request)
+    validate_request!(request)
+
+    result = dispatch_method(request["method"], request["params"] || {})
+
+    # JSON-RPC 2.0: notifications (requests without `id`) MUST NOT receive a response.
+    return nil unless request.key?("id")
+
+    McpToolkit::Protocol.success_response(id: request["id"], result:)
+  end
+
+  def validate_request!(request)
+    unless request["jsonrpc"] == McpToolkit::Protocol::JSONRPC_VERSION
+      raise McpToolkit::Protocol::InvalidRequest, "Missing jsonrpc version"
+    end
+    raise McpToolkit::Protocol::InvalidRequest, "Missing method" if request["method"].blank?
+  end
+
+  def dispatch_method(method, params)
+    case method
+    when "initialize"
+      handle_initialize(params)
+    when "initialized", "notifications/initialized"
+      handle_initialized
+    when "tools/list"
+      handle_tools_list
+    when "tools/call"
+      handle_tools_call(params)
+    when "ping"
+      handle_ping
+    else
+      return handle_upstream_list_changed(method) if upstream_list_changed_notification?(method)
+
+      raise McpToolkit::Protocol::MethodNotFound, method
+    end
+  end
+
+  # Satellites can tell the authority their tool list changed via a
+  # `notifications/<app>/tools/list_changed` notification, busting that upstream's
+  # cached aggregate. Matches the configured upstream keys only.
+  def upstream_list_changed_notification?(method)
+    upstream_key_from_notification(method).present?
+  end
+
+  def handle_upstream_list_changed(method)
+    McpToolkit::Gateway::Aggregator.new(config:).flush!(upstream_key_from_notification(method))
+    {}
+  end
+
+  def upstream_key_from_notification(method)
+    match = method.to_s.match(%r{\Anotifications/(?<key>.+)/tools/list_changed\z})
+    return nil unless match
+
+    config.upstreams.find(match[:key]) ? match[:key] : nil
+  end
+
+  def handle_initialize(params)
+    requested = params["protocolVersion"].to_s
+    versions = config.supported_protocol_versions
+    negotiated = versions.include?(requested) ? requested : versions.first
+
+    result = {
+      protocolVersion: negotiated,
+      capabilities: {
+        # listChanged: true — the aggregated list includes upstream tools, which
+        # can change when an upstream is reconfigured or sends a list_changed
+        # notification that busts the cached aggregate.
+        tools: { listChanged: true }
+      },
+      serverInfo: {
+        name: config.server_name,
+        version: config.server_version
+      }
+    }
+
+    # `instructions` is advertised on initialize only when configured, matching the
+    # SDK-backed satellite server (McpToolkit::Server.build) and the documented
+    # contract. Omitted when nil to keep the envelope clean.
+    result[:instructions] = config.server_instructions if config.server_instructions
+    result
+  end
+
+  def handle_initialized
+    {}
+  end
+
+  def handle_ping
+    {}
+  end
+
+  def handle_tools_list
+    {
+      tools: host_tool_definitions +
+        McpToolkit::Gateway::Aggregator.new(config:).tool_definitions(bearer_token: context.bearer_token)
+    }
+  end
+
+  # The host's own tool definitions, sourced from the injected tool_provider (the
+  # api-agnostic seam). `context` lets the provider hide superuser-only tools from
+  # a non-superuser caller. A host that registered no provider contributes none.
+  def host_tool_definitions
+    provider = config.tool_provider
+    return [] unless provider
+
+    provider.tool_definitions(context)
+  end
+
+  def handle_tools_call(params)
+    tool_name = params["name"]
+    arguments = params["arguments"] || {}
+
+    upstream = config.upstreams.split_tool_name(tool_name)
+    return handle_upstream_tools_call(upstream, arguments) if upstream
+
+    handle_host_tools_call(tool_name, arguments)
+  end
+
+  def handle_host_tools_call(tool_name, arguments)
+    tool = config.tool_provider&.find(tool_name)
+    raise McpToolkit::Protocol::MethodNotFound, "Tool not found: #{tool_name}" unless tool
+
+    ensure_tool_scope!(tool)
+
+    result = tool.call(context:, **symbolized_arguments(arguments))
+
+    {
+      content: [
+        {
+          type: "text",
+          text: result.is_a?(String) ? result : result.to_json
+        }
+      ]
+    }
+  end
+
+  # JSON gives string keys; a tool's `call(context:, **arguments)` needs symbol
+  # keys for the keyword splat. Deep-symbolized so nested argument hashes reach
+  # the tool in the same shape a symbol-keyed caller would pass.
+  def symbolized_arguments(arguments)
+    arguments.to_h.deep_symbolize_keys
+  end
+
+  def ensure_tool_scope!(tool)
+    required_scope = tool.required_permissions_scope
+    return if required_scope.blank?
+    return if context.principal&.authorized_for_scope?(required_scope)
+
+    raise McpToolkit::Protocol::InvalidRequest, "This token lacks the #{required_scope.inspect} scope"
+  end
+
+  def handle_upstream_tools_call((app_key, bare_tool_name), arguments)
+    McpToolkit::Gateway::Proxy.new(
+      app_key:,
+      tool_name: bare_tool_name,
+      account_id: context.account&.id,
+      bearer_token: context.bearer_token,
+      config:
+    ).call(arguments)
+  rescue McpToolkit::Gateway::UnknownUpstream => e
+    # The gateway stays transport-agnostic; the dispatcher maps an unknown
+    # upstream to its own "method not found".
+    raise McpToolkit::Protocol::MethodNotFound, e.message
+  rescue McpToolkit::Gateway::UpstreamCallError => e
+    raise translate_upstream_call_error(e)
+  end
+
+  # Translates the gateway's transport-agnostic upstream failure into the JSON-RPC
+  # error shape — verbatim relay of a satellite JSON-RPC error, else a generic
+  # internal error.
+  def translate_upstream_call_error(error)
+    if error.jsonrpc_error
+      McpToolkit::Protocol::Error.new(
+        error.jsonrpc_error["message"].to_s,
+        code: error.jsonrpc_error["code"] || McpToolkit::Protocol::ErrorCodes::INTERNAL_ERROR,
+        data: error.jsonrpc_error["data"]
+      )
+    else
+      McpToolkit::Protocol::InternalError.new(error.message)
+    end
+  end
+end

data/lib/mcp_toolkit/engine.rb

@@ -1,21 +1,36 @@
 # frozen_string_literal: true
 
-# Mountable Rails engine that draws the four MCP transport routes (defined in the
-# engine's config/routes.rb so they survive Rails' route reloads) against the
-# gem-provided McpToolkit::ServerController. A satellite mounts it in one line:
+# Mountable Rails engine that draws the MCP transport routes plus the authority
+# introspection route (defined in the engine's config/routes.rb so they survive
+# Rails' route reloads) against the gem-provided McpToolkit::ServerController /
+# McpToolkit::TokensController. A satellite mounts it in one line:
 #
 #   # config/routes.rb
 #   mount McpToolkit::Engine => "/mcp"
 #
 # yielding exactly the endpoints a hand-rolled satellite declared:
 #
-#   POST   /mcp          -> create   (JSON-RPC requests/responses)
-#   GET    /mcp          -> stream   (405; no server-initiated SSE)
-#   DELETE /mcp          -> destroy  (terminate the session)
-#   GET    /mcp/health   -> health   (unauthenticated probe)
+#   POST   /mcp                     -> create     (JSON-RPC requests/responses)
+#   GET    /mcp                     -> stream     (405; no server-initiated SSE)
+#   DELETE /mcp                     -> destroy    (terminate the session)
+#   GET    /mcp/health              -> health     (unauthenticated probe)
+#   POST   /mcp/tokens/introspect   -> introspect (authority token introspection;
+#                                                  drawn ONLY when auth_role is
+#                                                  :authority — a satellite that
+#                                                  mounts the engine gets no such
+#                                                  route)
 #
 # Loaded ONLY when Rails::Engine is available (see lib/mcp_toolkit.rb); the gem's
 # non-Rails consumers and its own unit suite never reference it.
 class McpToolkit::Engine < Rails::Engine
   isolate_namespace McpToolkit
+
+  # The gem-provided controllers subclass `config.parent_controller`, which the
+  # host sets in an initializer/to_prepare that must be READ AFTER it runs
+  # (Constraint B). They are therefore built lazily by
+  # `McpToolkit.build_engine_controllers!` (triggered via const_missing on first
+  # reference — at eager-load or first request). This resets them on every code
+  # reload so a changed parent (or a reloaded app parent class) takes effect on the
+  # next reference. Runs before `:eager_load!`, so the fresh classes exist for it.
+  config.to_prepare { McpToolkit.reset_engine_controllers! }
 end

data/lib/mcp_toolkit/engine_controllers.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+# Lazy `parent_controller` builder (Constraint B).
+#
+# The engine's controllers (McpToolkit::ServerController,
+# McpToolkit::TokensController) and the authority base
+# (McpToolkit::Authority::ServerController) all subclass
+# `config.parent_controller`. If that superclass were resolved in a class body of
+# an autoloaded/eager-loaded file, it could be read BEFORE the host's
+# initializer/to_prepare had set it — defaulting to ActionController::Base and, in
+# turn, breaking CSRF handling on the introspection endpoint.
+#
+# Instead, none of these controllers is a Zeitwerk-managed file. They are built
+# here from the CURRENT config, and the build is triggered LAZILY:
+#   * `const_missing` (below, and on McpToolkit::Authority) builds them the first
+#     time they are referenced — which, in a host, is at eager-load or first-
+#     request time, i.e. AFTER the app's initializers/to_prepare have run;
+#   * the engine's `config.to_prepare` RESETS them on every code reload so a
+#     changed `parent_controller` (or a reloaded app parent class) takes effect on
+#     the next reference.
+#
+# The whole `config/initializers/mcp_toolkit.rb` of a host can therefore live in
+# `to_prepare`: the parent is only ever read at build time.
+#
+# This file reopens `McpToolkit` to add module methods, so it is Zeitwerk-IGNORED
+# (like engine.rb) and required explicitly from the gem entry point.
+module McpToolkit
+  # The controllers built directly under McpToolkit (the engine's routes point at
+  # these). McpToolkit::Authority::ServerController is built alongside them but is
+  # fetched through McpToolkit::Authority's own const_missing.
+  ENGINE_CONTROLLER_NAMES = %i[ServerController TokensController].freeze
+
+  # (Re)builds the engine controllers + the authority base from the current
+  # config. Idempotent: an existing constant is replaced so a rebuild reflects a
+  # changed `parent_controller`. Reads `config.parent_controller` at call time.
+  def self.build_engine_controllers!
+    parent = config.parent_controller.constantize
+    define_controller(self, :ServerController, build_server_controller(parent))
+    define_controller(self, :TokensController, build_tokens_controller(parent))
+    define_controller(Authority, :ServerController, build_authority_server_controller(parent))
+    ServerController
+  end
+
+  # Undefines the built controllers so the next reference rebuilds them from the
+  # then-current config. Called from the engine's `to_prepare` on every reload.
+  def self.reset_engine_controllers!
+    [[self, :ServerController], [self, :TokensController], [Authority, :ServerController]].each do |mod, name|
+      mod.send(:remove_const, name) if mod.const_defined?(name, false)
+    end
+  end
+
+  # The SATELLITE transport controller the engine mounts (unchanged behavior; the
+  # SDK-backed path). Built lazily only so its parent is read after config.
+  def self.build_server_controller(parent)
+    Class.new(parent) { include McpToolkit::Transport::ControllerMethods }
+  end
+
+  # The AUTHORITY introspection endpoint the engine mounts at
+  # `POST /mcp/tokens/introspect`. Behavior is preserved exactly: it authenticates
+  # the bearer against `config.token_authenticator` (via Auth::Authority) and
+  # renders the introspection payload; a non-authority app answers `{ valid:
+  # false }` rather than erroring.
+  def self.build_tokens_controller(parent)
+    Class.new(parent) do
+      def introspect
+        token = McpToolkit::Auth::Authority.authenticate(mcp_extract_token, config: mcp_config)
+        return render(json: McpToolkit::Auth::Authority.invalid_payload, status: :unauthorized) if token.nil?
+
+        render json: McpToolkit::Auth::Authority.introspection_payload(token)
+      rescue McpToolkit::Errors::ConfigurationError
+        # Not configured as an authority (no token_authenticator): behave as if the
+        # token were invalid instead of surfacing a 500.
+        render json: McpToolkit::Auth::Authority.invalid_payload, status: :unauthorized
+      end
+
+      private
+
+      def mcp_config
+        McpToolkit.config
+      end
+
+      def mcp_extract_token
+        auth_header = request.headers["Authorization"]
+        return auth_header.sub("Bearer ", "") if auth_header&.start_with?("Bearer ")
+
+        request.headers["X-MCP-Token"].presence || params[:token].presence
+      end
+    end
+  end
+
+  # The AUTHORITY base controller a host subclasses (the recommended path for a
+  # host whose rate-limit/usage/account hooks touch app models).
+  def self.build_authority_server_controller(parent)
+    Class.new(parent) { include McpToolkit::Authority::ControllerMethods }
+  end
+
+  # Removes an existing same-named constant (avoiding a redefinition warning on a
+  # rebuild) before setting the freshly-built class.
+  def self.define_controller(mod, name, klass)
+    mod.send(:remove_const, name) if mod.const_defined?(name, false)
+    mod.const_set(name, klass)
+  end
+
+  # Backstop: build the engine controllers the first time one is referenced before
+  # any `to_prepare`/eager-load pass has built them (e.g. a bespoke route that
+  # names McpToolkit::ServerController directly). McpToolkit::Authority defines the
+  # sibling backstop for its ServerController.
+  def self.const_missing(name)
+    if ENGINE_CONTROLLER_NAMES.include?(name)
+      build_engine_controllers!
+      return const_get(name) if const_defined?(name, false)
+    end
+
+    super
+  end
+end

data/lib/mcp_toolkit/gateway/aggregator.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+# Aggregates the tool lists of all configured upstream MCP servers for a
+# gateway's `tools/list`, namespacing each tool as `<app>__<tool>`.
+#
+# Per upstream the namespaced list is cached in `config.cache_store`
+# (`config.upstream_list_ttl`, default 15 min) and bustable via `flush!` /
+# `flush!(key)`. On a cache miss the live pull runs per-upstream (one HTTP call
+# each); a failing/timeout upstream is omitted + logged, never breaking the list.
+#
+# Why the cache is safe to share globally: an upstream's `tools/list` is
+# token-INDEPENDENT — it returns the same public tool definitions to every valid
+# caller and enforces scope only when a tool is CALLED (per-call authorization is
+# the upstream's job). So one caller's pull is a correct answer for all callers.
+# This is a registration CONTRACT, not an assumption: an upstream that filters its
+# list by the caller's privilege (e.g. hides superuser-only tools) MUST register
+# `public_tool_list: false`, which opts it out of this cache (pulled live per
+# request) so a privileged caller's list can never be served to an unprivileged
+# one. See McpToolkit::Gateway::UpstreamRegistry::Upstream.
+#
+# Only a NON-EMPTY pull is ever cached. An empty or failed pull is almost always a
+# transient upstream hiccup (timeout, a session/handshake blip, a degenerate 200);
+# caching it would freeze a whole app's tools out of `tools/list` for the full TTL
+# for EVERY caller (a poisoned global cache), which is exactly the failure this
+# guards against. A stale empty already in the cache is treated as a miss and
+# re-pulled, so the aggregate self-heals as soon as the upstream returns tools.
+#
+# Concurrency: upstreams are pulled CONCURRENTLY via concurrent-ruby futures. When
+# running inside Rails, each future is wrapped in `Rails.application.executor` so
+# it participates in the framework's per-request lifecycle (reloading, query
+# cache, connection checkout); a non-Rails host runs plain futures. Output order
+# follows the registry order.
+class McpToolkit::Gateway::Aggregator
+  CACHE_KEY_PREFIX = "mcp_toolkit:gateway:tools:"
+
+  def initialize(config: McpToolkit.config)
+    @config = config
+  end
+
+  # Namespaced tool definitions across all configured upstreams. `bearer_token`
+  # is used only on a cache miss, so the upstream can authenticate the list
+  # request the same way it would a call. Upstreams are fetched concurrently;
+  # output order follows the registry order.
+  def tool_definitions(bearer_token: nil)
+    futures = config.upstreams.all.map do |upstream|
+      Concurrent::Promises.future do
+        within_executor { cached_or_live_definitions(upstream, bearer_token:) }
+      end
+    end
+
+    futures.flat_map(&:value!)
+  end
+
+  def flush!(key = nil)
+    if key
+      cache.delete(cache_key(key))
+    else
+      config.upstreams.all.each { |upstream| cache.delete(cache_key(upstream.key)) }
+    end
+  end
+
+  private
+
+  attr_reader :config
+
+  # Runs the block inside the Rails executor when a BOOTED Rails app is present
+  # (so a future participates in the request lifecycle: reloading, query cache,
+  # connection checkout), or plainly otherwise. Guards for `Rails` being defined
+  # but not booted (e.g. `rails/version` required without an initialized app), in
+  # which case `Rails.application` is nil and we run the block directly.
+  def within_executor(&)
+    if defined?(Rails) && Rails.respond_to?(:application) && Rails.application
+      Rails.application.executor.wrap(&)
+    else
+      yield
+    end
+  end
+
+  def cached_or_live_definitions(upstream, bearer_token:)
+    # A caller-dependent upstream opts out of the shared cache (keyed by upstream
+    # only): caching it would leak one caller's list to callers of other privilege.
+    return live_definitions(upstream, bearer_token:) unless upstream.public_tool_list
+
+    cached = cache.read(cache_key(upstream.key))
+    # `present?` treats both a nil miss AND a stale empty list as "no cache", so a
+    # previously poisoned empty entry is re-pulled instead of served.
+    return cached if cached.present?
+
+    definitions = live_definitions(upstream, bearer_token:)
+    # Only persist a real, non-empty list; never cache an empty/degraded pull.
+    cache.write(cache_key(upstream.key), definitions, expires_in: config.upstream_list_ttl) if definitions.present?
+    definitions
+  rescue McpToolkit::Gateway::Client::Error => e
+    # Degrade gracefully: omit this upstream's tools, don't cache the failure.
+    config.logger&.error("MCP upstream #{upstream.key} tools/list failed, omitting: #{e.message}")
+    []
+  end
+
+  def live_definitions(upstream, bearer_token:)
+    client = McpToolkit::Gateway::Client.new(upstream:, bearer_token:, config:)
+    client.tools_list.map do |definition|
+      namespaced(upstream, definition)
+    end
+  end
+
+  # Re-keys an upstream tool definition into the gateway's aggregate namespace.
+  def namespaced(upstream, definition)
+    definition = definition.dup
+    definition["name"] = upstream.name_for(definition["name"])
+    definition
+  end
+
+  def cache
+    config.cache_store
+  end
+
+  def cache_key(key)
+    "#{CACHE_KEY_PREFIX}#{key}"
+  end
+end