mcp_toolkit 0.3.0 → 0.5.0

data/lib/mcp_toolkit/gateway/client.rb

@@ -0,0 +1,305 @@
+# frozen_string_literal: true
+
+require "faraday"
+
+# Minimal MCP client over Streamable HTTP, used by a GATEWAY to talk to an
+# upstream MCP server when aggregating its tool list and proxying tool calls.
+#
+# It speaks the same Streamable-HTTP MCP that McpToolkit::Transport serves:
+#   1. POST `initialize`                 -> capture the `Mcp-Session-Id` response header
+#   2. POST `notifications/initialized`  (a notification; no response expected)
+#   3. POST `tools/list` / `tools/call`, echoing the session header
+#
+# Auth & account selection are pass-through: the caller's bearer token and the
+# account selector (`_meta`) are forwarded so the upstream can introspect the
+# same token against its authority and resolve the same account.
+#
+# Transport notes
+#   - Content negotiation: we POST application/json and Accept both
+#     application/json and text/event-stream. If the upstream answers with SSE
+#     (one message + EOF, as the toolkit's own transport does) we extract the
+#     single JSON payload from the `data:` line.
+#   - Every public method may raise McpToolkit::Gateway::Client::Error (timeouts,
+#     non-2xx, unparseable bodies, JSON-RPC errors). Callers decide whether to
+#     degrade (omit from list) or surface the error (proxied call).
+#
+# Everything app-specific (server identity, protocol version, timeout, logger) is
+# injected via McpToolkit::Configuration; nothing here names a deployment.
+class McpToolkit::Gateway::Client
+  SESSION_HEADER = "Mcp-Session-Id"
+  JSONRPC_VERSION = "2.0"
+
+  # The protocol version offered on the handshake when the config does not pin
+  # one. Sourced from the wrapped `mcp` SDK's latest-supported constant when
+  # available, with a literal fallback so the gem still loads if the SDK moves the
+  # constant. `config.protocol_version` overrides it.
+  DEFAULT_PROTOCOL_VERSION =
+    if defined?(MCP::Configuration) && MCP::Configuration.const_defined?(:LATEST_STABLE_PROTOCOL_VERSION)
+      MCP::Configuration::LATEST_STABLE_PROTOCOL_VERSION
+    else
+      "2025-06-18"
+    end
+
+  # An upstream whose session store is not shared across its pods can answer a
+  # NON-initialize request with HTTP 404 (the pod fielding it never saw our
+  # initialize), or with a JSON-RPC -32001 "Session not found or expired". Both
+  # mean the same thing: our session is gone and a fresh handshake recovers it.
+  SESSION_LOSS_HTTP_STATUS = 404
+  SESSION_NOT_FOUND_CODE = -32_001
+
+  # Raised for any upstream failure. `jsonrpc_error` carries an upstream JSON-RPC
+  # error hash when the failure was a protocol-level error response, so a proxy
+  # can relay it verbatim; nil for transport/HTTP failures. `http_status` carries
+  # the HTTP status for a non-2xx response (nil otherwise), so callers can
+  # distinguish a session-loss 404 from other failures. This class references NO
+  # transport/protocol-error type — the consumer maps it.
+  class Error < McpToolkit::Error
+    attr_reader :jsonrpc_error, :http_status
+
+    def initialize(message, jsonrpc_error: nil, http_status: nil)
+      super(message)
+      @jsonrpc_error = jsonrpc_error
+      @http_status = http_status
+    end
+  end
+
+  attr_reader :upstream, :bearer_token, :timeout, :config
+
+  def initialize(upstream:, bearer_token: nil, config: McpToolkit.config)
+    @upstream = upstream
+    @bearer_token = bearer_token
+    @config = config
+    @timeout = config.upstream_timeout
+    @session_id = nil
+    @initialized = false
+  end
+
+  # Returns the upstream's raw tools array (each a tool definition hash with
+  # string keys: "name", "description", "inputSchema"). Bare names — the caller
+  # namespaces them.
+  def tools_list
+    with_session_recovery("tools/list") do
+      result = rpc!("tools/list")
+      Array(result["tools"])
+    end
+  end
+
+  # Proxies a tools/call to the upstream. `arguments` and `meta` are forwarded
+  # as-is. Returns the upstream's `result` hash verbatim (typically
+  # `{ "content" => [...] }`). A JSON-RPC error from the upstream is raised as an
+  # Error carrying `jsonrpc_error` so a proxy can relay it.
+  def tools_call(name:, arguments: {}, meta: nil)
+    params = { "name" => name, "arguments" => arguments }
+    params["_meta"] = meta if meta.present?
+    with_session_recovery("tools/call #{name}") do
+      rpc!("tools/call", params)
+    end
+  end
+
+  private
+
+  # The protocol version to offer on the handshake: the config's pin, else the
+  # gem's default.
+  def protocol_version
+    config.protocol_version || DEFAULT_PROTOCOL_VERSION
+  end
+
+  # Runs a request-bearing RPC against an initialized session, transparently
+  # recovering from a SINGLE session-loss response: if the upstream reports our
+  # session is gone (see SESSION_LOSS_HTTP_STATUS / SESSION_NOT_FOUND_CODE) we
+  # drop the dead session, re-handshake, and retry the block exactly ONCE. Any
+  # other error — a genuine JSON-RPC tool error, a timeout, a non-404 HTTP
+  # failure — propagates verbatim, unchanged. An upstream tool that returns an
+  # `isError` content result is a normal return value here (not an exception), so
+  # it passes straight through.
+  def with_session_recovery(operation, &)
+    ensure_initialized!
+    yield
+  rescue Error => e
+    raise unless session_loss?(e)
+
+    recover_and_retry(operation, e, &)
+  end
+
+  # Second (and final) attempt after a session-loss signal: re-establish the
+  # session and run the block once more. Bounded to one retry — this method does
+  # not recurse into `with_session_recovery`, so it cannot loop. If re-init or the
+  # retry itself fails, the error is clarified before it propagates.
+  def recover_and_retry(operation, original_error)
+    log_session_recovery(operation, original_error)
+    reset_session!
+    ensure_initialized!
+    yield
+  rescue Error => e
+    raise clarified_session_error(e, operation)
+  end
+
+  # Performs the initialize handshake once per client instance, capturing the
+  # session id, then sends the `notifications/initialized` notification.
+  def ensure_initialized!
+    return if @initialized
+
+    response = post_jsonrpc(
+      jsonrpc_request("initialize", {
+                        "protocolVersion" => protocol_version,
+                        "capabilities" => {},
+                        # The GATEWAY-client identity (falls back to the server
+                        # identity), split so an authority can advertise its own
+                        # server_name to callers while keeping this upstream
+                        # handshake byte-identical.
+                        "clientInfo" => {
+                          "name" => config.gateway_client_name,
+                          "version" => config.gateway_client_version
+                        }
+                      })
+    )
+    @session_id = response.headers[SESSION_HEADER].presence
+    parse_rpc_result!(response)
+
+    # Best-effort lifecycle notification; upstreams may ignore it. A notification
+    # returns no body (202), so we don't parse a result.
+    post_jsonrpc(jsonrpc_notification("notifications/initialized"))
+
+    @initialized = true
+  end
+
+  # Drops the current session so the next `ensure_initialized!` re-handshakes.
+  def reset_session!
+    @initialized = false
+    @session_id = nil
+  end
+
+  # True when an upstream error signals our session is gone and a fresh handshake
+  # could recover it: a bare HTTP 404, or a JSON-RPC -32001 session error.
+  def session_loss?(error)
+    error.http_status == SESSION_LOSS_HTTP_STATUS ||
+      error.jsonrpc_error&.dig("code") == SESSION_NOT_FOUND_CODE
+  end
+
+  # A session-loss failure we could NOT recover from (already retried once).
+  # Preserve verbatim relay for a genuine JSON-RPC error and pass through any
+  # unrelated failure; only rewrite the bare "returned HTTP 404", whose message
+  # hides the real cause.
+  def clarified_session_error(error, operation)
+    return error if error.jsonrpc_error
+    return error unless session_loss?(error)
+
+    Error.new(
+      "upstream #{upstream.key} #{operation} failed: the MCP session was lost and could not " \
+      "be re-established after re-initializing the handshake",
+      http_status: error.http_status
+    )
+  end
+
+  def log_session_recovery(operation, error)
+    config.logger&.warn(
+      "MCP upstream #{upstream.key} #{operation}: session lost (#{session_loss_reason(error)}), " \
+      "re-initializing and retrying once"
+    )
+  end
+
+  def session_loss_reason(error)
+    return "JSON-RPC #{error.jsonrpc_error["code"]}" if error.jsonrpc_error
+    return "HTTP #{error.http_status}" if error.http_status
+
+    "unknown"
+  end
+
+  # Sends a request-bearing JSON-RPC call and returns its `result`.
+  def rpc!(method, params = {})
+    response = post_jsonrpc(jsonrpc_request(method, params))
+    parse_rpc_result!(response)
+  end
+
+  def jsonrpc_request(method, params = {})
+    { "jsonrpc" => JSONRPC_VERSION, "id" => SecureRandom.uuid, "method" => method, "params" => params }
+  end
+
+  def jsonrpc_notification(method, params = {})
+    { "jsonrpc" => JSONRPC_VERSION, "method" => method, "params" => params }
+  end
+
+  def post_jsonrpc(body)
+    connection.post(upstream.url) do |request|
+      apply_request_headers(request.headers)
+      request.body = JSON.generate(body)
+    end
+  rescue Faraday::TimeoutError => e
+    raise upstream_error("timed out after #{timeout}s", e)
+  rescue Faraday::Error => e
+    raise upstream_error("request failed", e)
+  end
+
+  def upstream_error(reason, cause)
+    Error.new("upstream #{upstream.key} #{reason}: #{cause.message}")
+  end
+
+  # Sets the content-negotiation, auth, and session headers on an outgoing
+  # request. Auth and session headers are added only when present.
+  def apply_request_headers(headers)
+    headers["Content-Type"] = "application/json"
+    headers["Accept"] = "application/json, text/event-stream"
+    headers["Authorization"] = "Bearer #{bearer_token}" if bearer_token.present?
+    headers[SESSION_HEADER] = @session_id if @session_id.present?
+  end
+
+  # Validates the HTTP response and extracts the JSON-RPC `result`, raising on a
+  # JSON-RPC error (with the error hash attached for verbatim relay).
+  def parse_rpc_result!(response)
+    raise_http_error!(response) unless response.success?
+
+    payload = decode_body(response)
+    return {} if payload.nil? # e.g. 202 Accepted for a notification
+
+    if payload["error"]
+      raise Error.new(
+        "upstream #{upstream.key} JSON-RPC error: #{payload.dig("error", "message")}",
+        jsonrpc_error: payload["error"]
+      )
+    end
+
+    payload["result"] || {}
+  end
+
+  # The `http_status` is carried on the Error so a session-loss 404 can be told
+  # apart from other non-2xx failures (see `session_loss?`).
+  def raise_http_error!(response)
+    raise Error.new("upstream #{upstream.key} returned HTTP #{response.status}", http_status: response.status)
+  end
+
+  # Decodes either a plain JSON body or a single SSE `data:` frame.
+  def decode_body(response)
+    json = json_payload(response)
+    return nil if json.nil?
+
+    JSON.parse(json)
+  rescue JSON::ParserError => e
+    raise Error, "upstream #{upstream.key} returned unparseable body: #{e.message}"
+  end
+
+  # Extracts the JSON text to parse, unwrapping an SSE frame when the response is
+  # an event stream. Returns nil for an empty body or an empty/absent frame.
+  def json_payload(response)
+    body = response.body.to_s
+    return nil if body.strip.empty?
+
+    content_type = response.headers["Content-Type"].to_s
+    json = content_type.include?("text/event-stream") ? extract_sse_data(body) : body
+    json unless json.nil? || json.strip.empty?
+  end
+
+  # Pulls the JSON payload from the first `data:` line of an SSE stream.
+  def extract_sse_data(body)
+    body.each_line.filter_map do |line|
+      line.start_with?("data:") ? line.sub(/\Adata:\s?/, "").chomp : nil
+    end.first
+  end
+
+  def connection
+    @connection ||= Faraday.new do |conn|
+      conn.options.timeout = timeout
+      conn.options.open_timeout = timeout
+      conn.adapter Faraday.default_adapter
+    end
+  end
+end

data/lib/mcp_toolkit/gateway/proxy.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+# Proxies a namespaced upstream tool call (`<app>__<tool>`) to its upstream MCP
+# server.
+#
+# The account selector is forwarded as `_meta[config.account_meta_key]` when an
+# `account_id` is supplied. The caller passes the already-resolved account id (a
+# scalar), not a domain object — resolving the tenant is the consumer's job.
+#
+# Error mapping is deliberately thin and transport-agnostic:
+#   * an unregistered `app_key` raises McpToolkit::Gateway::UnknownUpstream;
+#   * an upstream call failure is re-raised as McpToolkit::Gateway::UpstreamCallError
+#     carrying the upstream's `jsonrpc_error` / `http_status`.
+# Neither is mapped to a JSON-RPC/protocol error class here — the consuming
+# dispatcher does that at its call site.
+class McpToolkit::Gateway::Proxy
+  attr_reader :app_key, :tool_name, :account_id, :bearer_token, :config
+
+  def initialize(app_key:, tool_name:, account_id: nil, bearer_token: nil, config: McpToolkit.config)
+    @app_key = app_key
+    @tool_name = tool_name
+    @account_id = account_id
+    @bearer_token = bearer_token
+    @config = config
+  end
+
+  def call(arguments)
+    upstream = config.upstreams.find(app_key)
+    raise McpToolkit::Gateway::UnknownUpstream, "Unknown application: #{app_key}" if upstream.nil?
+
+    client = McpToolkit::Gateway::Client.new(upstream:, bearer_token:, config:)
+    client.tools_call(name: tool_name, arguments:, meta:)
+  rescue McpToolkit::Gateway::Client::Error => e
+    relay_upstream_error(e)
+  end
+
+  private
+
+  def relay_upstream_error(error)
+    log_proxied_failure(error)
+
+    raise McpToolkit::Gateway::UpstreamCallError.new(
+      error.message,
+      jsonrpc_error: error.jsonrpc_error,
+      http_status: error.http_status
+    )
+  end
+
+  # Emit one concise, greppable ERROR line per failed proxied call — never a
+  # bearer token or a full response body.
+  def log_proxied_failure(error)
+    config.logger&.error(
+      "MCP upstream #{app_key} tools/call #{tool_name} failed#{failure_detail(error)}: #{error.message}"
+    )
+  end
+
+  def failure_detail(error)
+    code = error.jsonrpc_error&.dig("code")
+    return " (jsonrpc_code=#{code})" if code
+    return " (http_status=#{error.http_status})" if error.http_status
+
+    ""
+  end
+
+  def meta
+    return nil if account_id.nil?
+
+    { config.account_meta_key => account_id }
+  end
+end

data/lib/mcp_toolkit/gateway/unknown_upstream.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+# Raised by McpToolkit::Gateway::Proxy when a namespaced tool call targets an
+# upstream key that is not registered. The GEM does NOT translate this into any
+# JSON-RPC / protocol error class — the consuming dispatcher maps it to whatever
+# error shape its transport speaks (e.g. a "method not found"). Kept a plain
+# McpToolkit::Error so the gem stays transport-agnostic.
+class McpToolkit::Gateway::UnknownUpstream < McpToolkit::Error; end

data/lib/mcp_toolkit/gateway/upstream_call_error.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# Raised by McpToolkit::Gateway::Proxy when a proxied `tools/call` fails at the
+# upstream. It carries the upstream failure detail so the CONSUMER can decide how
+# to surface it:
+#
+#   * `jsonrpc_error` — the upstream's JSON-RPC error hash when the failure was a
+#     protocol-level error response (so the consumer can relay it verbatim); nil
+#     for transport/HTTP failures.
+#   * `http_status`   — the HTTP status for a non-2xx response (nil otherwise).
+#
+# The gem deliberately does NOT map this to a protocol/transport error class:
+# that mapping lives in the consuming dispatcher, keeping the gateway
+# transport-agnostic.
+class McpToolkit::Gateway::UpstreamCallError < McpToolkit::Error
+  attr_reader :jsonrpc_error, :http_status
+
+  def initialize(message, jsonrpc_error: nil, http_status: nil)
+    super(message)
+    @jsonrpc_error = jsonrpc_error
+    @http_status = http_status
+  end
+end

data/lib/mcp_toolkit/gateway/upstream_registry.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+# Registry of upstream MCP servers a GATEWAY app aggregates and proxies to.
+#
+# Each upstream has a `key` (the tool-name namespace prefix — its tools are
+# surfaced as `<key>__<tool>`) and a `url` (the upstream's MCP HTTP endpoint).
+# Upstreams are registered at boot (typically from ENV); an upstream whose url is
+# blank is never registered, so an unconfigured environment behaves exactly like
+# a gateway with no upstreams.
+#
+# Unlike a global module singleton, this is a PER-CONFIG instance (like
+# `config.registry`): each `McpToolkit::Configuration` carries its own, exposed as
+# `config.upstreams`. That gives test isolation for free (a fresh config resets
+# it) and matches the gem's per-config convention. Register via the config sugar
+# `config.register_upstream(key:, url:)` or directly on `config.upstreams`.
+class McpToolkit::Gateway::UpstreamRegistry
+  # Separates an app key from a tool name in an aggregated tool name. A double
+  # underscore so it doesn't collide with single underscores in bare tool names
+  # (e.g. "list_items"); it also matches the gem's `<app>__<action>` scope
+  # separator.
+  NAMESPACE_SEPARATOR = "__"
+
+  # `public_tool_list` is the gateway-cache contract: an upstream's `tools/list`
+  # MUST be caller-INDEPENDENT (the same public tools for every valid token; scope
+  # is enforced only when a tool is CALLED). That invariant is what lets the
+  # Aggregator share one cache entry per upstream across all callers. An upstream
+  # that filters its list by the caller's privilege (e.g. hides superuser-only
+  # tools) breaks it and MUST register `public_tool_list: false`, which opts it out
+  # of the shared cache so a privileged caller's list can't leak to an
+  # unprivileged one. Default true — the overwhelming common case.
+  Upstream = Data.define(:key, :url, :public_tool_list) do
+    def initialize(key:, url:, public_tool_list: true)
+      super
+    end
+
+    def name_for(tool_name)
+      "#{key}#{NAMESPACE_SEPARATOR}#{tool_name}"
+    end
+  end
+
+  def initialize
+    @registered = {}
+  end
+
+  # Registers an upstream by key. A blank url is ignored, so callers can pass an
+  # ENV lookup directly without guarding it. Pass `public_tool_list: false` for an
+  # upstream whose tool list varies by caller privilege, opting it out of the
+  # shared list cache (see the Upstream contract above).
+  def register(key:, url:, public_tool_list: true)
+    return if url.blank?
+
+    @registered[key.to_s] = Upstream.new(key: key.to_s, url:, public_tool_list:)
+  end
+
+  # Clears every registered upstream.
+  def reset!
+    @registered = {}
+  end
+
+  # All registered upstreams (insertion order).
+  def all
+    @registered.values
+  end
+
+  # The registered upstream for a key, or nil.
+  def find(key)
+    @registered[key.to_s]
+  end
+
+  # `<app>__<tool>` -> [app_key, bare_tool_name] for a registered upstream; nil
+  # for an un-namespaced name or an unknown/unregistered key.
+  def split_tool_name(tool_name)
+    prefix, separator, rest = tool_name.to_s.partition(NAMESPACE_SEPARATOR)
+    return nil if separator.empty? || rest.empty?
+    return nil unless find(prefix)
+
+    [prefix, rest]
+  end
+end

data/lib/mcp_toolkit/list_executor.rb

@@ -45,14 +45,41 @@ class McpToolkit::ListExecutor
     relation = resource.resolve_relation(scope_root)
     relation = apply_ids(relation)
     relation = apply_updated_since(relation)
+    relation = apply_custom_filters(relation)
     relation = apply_attribute_filters(relation)
     apply_order(relation)
   end
 
-  # Order by `id` when the primary key is numeric; otherwise by `created_at`
-  # (a non-numeric PK does not sort meaningfully).
+  # Applies the resource's declared custom (resource-specific) filters — each an
+  # arbitrary host-supplied block — for the keys actually present as TOP-LEVEL
+  # request params, BEFORE the generic allowlist `filter` attributes. Each block
+  # receives the already-scoped relation and the raw value and returns a narrowed
+  # relation, so a host can express a relational filter the generic path can't
+  # derive (see Resource#filter). A resource with no custom filters is a no-op.
+  def apply_custom_filters(relation)
+    resource.custom_filters.each_value do |custom_filter|
+      value = params[custom_filter.name]
+      next if value.nil? || value == ""
+
+      relation = custom_filter.applier.call(relation, value)
+    end
+    relation
+  end
+
+  # Order by `id` when the primary key is numeric; otherwise per
+  # `config.non_numeric_pk_order`: `created_at` with the primary key as a
+  # tiebreaker (default — rows bulk-inserted in one transaction share a
+  # `created_at`, and without a total order offset pagination could duplicate
+  # or skip rows), or the primary key alone for a host preserving a pre-gem
+  # order-by-id contract. `numeric_primary_key?` returning false guarantees the
+  # model exposes a non-nil primary key, so it can be read directly here.
   def apply_order(relation)
-    relation.order(numeric_primary_key? ? :id : :created_at)
+    return relation.order(:id) if numeric_primary_key?
+
+    pk = resource.model.primary_key.to_sym
+    return relation.order(pk) if McpToolkit.config.non_numeric_pk_order == :primary_key
+
+    relation.order(:created_at, pk)
   end
 
   def numeric_primary_key?
@@ -81,9 +108,10 @@ class McpToolkit::ListExecutor
 
     mapping = resource.filterable_columns
     validate_filter_keys!(filter, mapping)
+    validate_filter_companions!(filter)
 
     filter.each do |request_key, value|
-      next if value.nil? || value == ""
+      next if skipped_filter_value?(value)
 
       column = mapping[request_key.to_sym]
       relation = McpToolkit::Filtering.apply(relation, column, value)
@@ -91,6 +119,35 @@ class McpToolkit::ListExecutor
     relation
   end
 
+  # Under :tokenized semantics an empty string means "no filter" (a JSON null
+  # still flows through as an IS NULL filter, like the "null" token — see
+  # McpToolkit::Filtering); under :literal every value reaches the WHERE clause
+  # verbatim.
+  def skipped_filter_value?(value)
+    value == "" && McpToolkit.config.bare_filter_value_semantics != :literal
+  end
+
+  # A filter key may declare a companion key it cannot be used without (e.g. a
+  # polymorphic foreign key is type-ambiguous without its `*_type`) — see
+  # Resource#filter_requirements. Rejected up front rather than producing a
+  # subtly wrong WHERE. A key whose value the apply loop will SKIP counts as
+  # not provided — a skipped ("" under :tokenized) companion must not satisfy
+  # the requirement, or the FK would be applied alone: exactly the
+  # type-ambiguous WHERE the requirement exists to prevent.
+  def validate_filter_companions!(filter)
+    requirements = resource.filter_requirements
+    return if requirements.empty?
+
+    provided = filter.reject { |_key, value| skipped_filter_value?(value) }.keys.map(&:to_sym)
+    requirements.each do |key, required|
+      next unless provided.include?(key)
+      next if provided.include?(required)
+
+      raise McpToolkit::Errors::InvalidParams,
+            "filter attribute #{key} requires #{required} to also be provided"
+    end
+  end
+
   def validate_filter_keys!(filter, mapping)
     keys = filter.keys.map(&:to_sym)
     unknown = keys - mapping.keys
@@ -106,6 +163,10 @@ class McpToolkit::ListExecutor
     return relation if params[:ids].blank?
 
     ids = params[:ids].to_s.split(",").map(&:strip).compact_blank
+    # Bound the IN-set the same way the per-attribute filters are bounded — the
+    # top-level `ids` param builds `WHERE id IN (...)` on its own path, so it
+    # must honor config.max_filter_values too (else it's an unbounded IN clause).
+    McpToolkit::Filtering.enforce_filter_limit!(ids.length, McpToolkit.config)
     ids.empty? ? relation : relation.where(id: ids)
   end