omni_agent 0.1.9 → 0.1.11

data/lib/omni_agent/mcp/oauth/access_token.rb

@@ -0,0 +1,60 @@
+module OmniAgent
+  module MCP
+    module OAuth
+      # A validated access token and its claims (JWT payload or introspection
+      # response). It is the principal unless `authenticate :oauth` has a block.
+      class AccessToken
+        attr_reader :token, :claims
+
+        def initialize(token, claims)
+          @token = token
+          @claims = claims.transform_keys(&:to_s).freeze
+        end
+
+        def [](key)
+          claims[key.to_s]
+        end
+
+        def subject
+          claims["sub"]
+        end
+
+        def issuer
+          claims["iss"]
+        end
+
+        def client_id
+          claims["client_id"] || claims["azp"]
+        end
+
+        def audience
+          Array(claims["aud"]).map(&:to_s)
+        end
+
+        # RFC 9068 / RFC 7662 use a space-separated "scope"; some providers
+        # (Okta, Entra ID) use "scp", as a string or an array.
+        def scopes
+          raw = claims.key?("scope") ? claims["scope"] : claims["scp"]
+          raw.is_a?(Array) ? raw.map(&:to_s) : raw.to_s.split
+        end
+
+        def scope?(*names)
+          (names.flatten.map(&:to_s) - scopes).empty?
+        end
+
+        def expires_at
+          Time.at(claims["exp"]) if claims["exp"].is_a?(Numeric)
+        end
+
+        def to_s
+          subject.to_s
+        end
+
+        # Never print the raw token.
+        def inspect
+          "#<#{self.class.name} sub=#{subject.inspect} client_id=#{client_id.inspect} scopes=#{scopes.inspect}>"
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/http_client.rb

@@ -0,0 +1,52 @@
+require "json"
+require "net/http"
+require "openssl"
+require "uri"
+
+module OmniAgent
+  module MCP
+    module OAuth
+      # Minimal JSON client for the authorization server (JWKS, introspection).
+      module HTTPClient
+        OPEN_TIMEOUT = 5
+        READ_TIMEOUT = 5
+
+        NETWORK_ERRORS = [
+          IOError, SystemCallError, SocketError, Timeout::Error, OpenSSL::SSL::SSLError, URI::InvalidURIError
+        ].freeze
+
+        module_function
+
+        def get_json(url)
+          uri = URI.parse(url.to_s)
+          request(uri, Net::HTTP::Get.new(uri))
+        end
+
+        def post_form_json(url, form, basic_auth: nil)
+          uri = URI.parse(url.to_s)
+          post = Net::HTTP::Post.new(uri)
+          post.set_form_data(form)
+          post.basic_auth(*basic_auth) if basic_auth
+          request(uri, post)
+        end
+
+        def request(uri, http_request)
+          http_request["accept"] = "application/json"
+
+          response = Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https",
+                                     open_timeout: OPEN_TIMEOUT, read_timeout: READ_TIMEOUT) do |http|
+            http.request(http_request)
+          end
+
+          unless response.is_a?(Net::HTTPSuccess)
+            raise VerificationError, "#{http_request.method} #{uri} returned HTTP #{response.code}"
+          end
+
+          JSON.parse(response.body.to_s)
+        rescue JSON::ParserError, *NETWORK_ERRORS => e
+          raise VerificationError, "#{http_request.method} #{uri} failed: #{e.class}: #{e.message}"
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/introspection_verifier.rb

@@ -0,0 +1,43 @@
+require "uri"
+
+module OmniAgent
+  module MCP
+    module OAuth
+      # Asks the authorization server whether a token is active (RFC 7662).
+      # Works with opaque tokens. Returns the introspection response, or nil
+      # when the token is not active.
+      class IntrospectionVerifier
+        def initialize(url:, client_id: nil, client_secret: nil)
+          raise ArgumentError, "introspection: requires url:" if url.nil?
+
+          @url = url
+          @client_id = client_id
+          @client_secret = client_secret
+        end
+
+        def call(token, _request = nil)
+          response = HTTPClient.post_form_json(
+            OAuth.resolve(@url),
+            { "token" => token, "token_type_hint" => "access_token" },
+            basic_auth: credentials
+          )
+          response.is_a?(Hash) && response["active"] == true ? response : nil
+        end
+
+        def inspect
+          "#<#{self.class.name} url=#{@url.inspect}>"
+        end
+
+        private
+
+        # client_secret_basic: RFC 6749 section 2.3.1 form-encodes both values.
+        def credentials
+          client_id = OAuth.resolve(@client_id)
+          return if client_id.nil?
+
+          [ client_id, OAuth.resolve(@client_secret) ].map { |value| URI.encode_www_form_component(value.to_s) }
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/jwt_verifier.rb

@@ -0,0 +1,122 @@
+module OmniAgent
+  module MCP
+    module OAuth
+      # Verifies JWT access tokens (RFC 9068) with the authorization server's
+      # JSON Web Key Set. Returns the payload, or nil when the token is invalid.
+      # Issuer, audience and scopes are checked afterwards by ResourceServer.
+      class JWTVerifier
+        # Asymmetric algorithms only: a JWKS publishes public keys.
+        DEFAULT_ALGORITHMS = %w[RS256 RS384 RS512 PS256 PS384 PS512 ES256 ES384 ES512].freeze
+        CACHE_TTL = 300
+        # A token signed with an unknown `kid` refetches the key set (key
+        # rotation), but at most this often, so forged tokens cannot make every
+        # request call the authorization server. Also the wait after a failed fetch.
+        MIN_REFRESH_INTERVAL = 30
+
+        def self.require_jwt!
+          require "jwt"
+        rescue LoadError
+          raise OmniAgent::MissingDependencyError,
+                "The 'jwt' gem is required for `authenticate :oauth, jwks_uri: ...`. " \
+                "Please add `gem 'jwt'` to your Gemfile."
+        end
+
+        def initialize(jwks_uri, algorithms: DEFAULT_ALGORITHMS, leeway: 0)
+          self.class.require_jwt!
+
+          @jwks_uri = jwks_uri
+          @algorithms = Array(algorithms).map(&:to_s)
+          @leeway = leeway
+          # Held only while fetching, so at most one fetch runs at a time.
+          # Reading the cached keys never takes it.
+          @fetch_mutex = Mutex.new
+          @jwks = nil
+          @fetched_at = nil
+          @failed_at = nil
+        end
+
+        def call(token, _request = nil)
+          payload, _header = ::JWT.decode(
+            token, nil, true,
+            algorithms: @algorithms,
+            jwks: method(:jwks),
+            required_claims: %w[exp iss],
+            leeway: @leeway
+          )
+          payload
+        rescue ::JWT::DecodeError
+          nil
+        end
+
+        private
+
+        # Key loader for JWT.decode. Called with `kid_not_found: true` when the
+        # token's `kid` is not in the cached set.
+        def jwks(options = {})
+          if @jwks.nil?
+            load_initial_jwks
+          elsif refresh_due?(options) && @fetch_mutex.try_lock
+            # Only the thread that got the lock refreshes; the others keep
+            # using the cached keys instead of waiting for it.
+            begin
+              fetch_jwks
+            ensure
+              @fetch_mutex.unlock
+            end
+          end
+
+          @jwks
+        end
+
+        # With no keys yet, concurrent requests wait for one shared fetch rather
+        # than each fetching in turn. After a failure they fail fast until
+        # MIN_REFRESH_INTERVAL has passed, so a down authorization server does
+        # not tie up every app thread on timeouts.
+        def load_initial_jwks
+          raise_unavailable if backing_off?
+
+          @fetch_mutex.synchronize do
+            next if @jwks
+            raise_unavailable if backing_off?
+
+            fetch_jwks
+          end
+
+          raise_unavailable if @jwks.nil?
+        end
+
+        def refresh_due?(options)
+          return false if backing_off?
+
+          age = monotonic_now - @fetched_at
+          age > CACHE_TTL || (options[:kid_not_found] && age > MIN_REFRESH_INTERVAL)
+        end
+
+        def fetch_jwks
+          keys = HTTPClient.get_json(OAuth.resolve(@jwks_uri))
+          # @fetched_at first: readers that see @jwks expect it to be set.
+          @fetched_at = monotonic_now
+          @jwks = keys
+          @failed_at = nil
+        rescue VerificationError => e
+          @failed_at = monotonic_now
+          fallback = @jwks ? "using cached keys" : "retrying in #{MIN_REFRESH_INTERVAL}s"
+          MCP.logger&.warn("[OmniAgent::MCP] Could not fetch JWKS, #{fallback}: #{e.message}")
+        end
+
+        def backing_off?
+          failed_at = @failed_at
+          !failed_at.nil? && monotonic_now - failed_at < MIN_REFRESH_INTERVAL
+        end
+
+        def raise_unavailable
+          raise VerificationError, "JWKS from #{OAuth.resolve(@jwks_uri)} is unavailable"
+        end
+
+        def monotonic_now
+          Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth/resource_server.rb

@@ -0,0 +1,166 @@
+require "uri"
+
+module OmniAgent
+  module MCP
+    module OAuth
+      # Built by `authenticate :oauth, ...`. Validates bearer access tokens
+      # (signature or introspection, issuer, audience, expiry and scopes) and
+      # describes the server in Protected Resource Metadata (RFC 9728) so MCP
+      # clients can find the authorization server.
+      class ResourceServer
+        METADATA_PATH = "/.well-known/oauth-protected-resource".freeze
+        # Seconds of clock skew tolerated when checking exp and nbf.
+        LEEWAY = 30
+
+        attr_reader :scopes
+
+        def initialize(resource:, authorization_servers:, scopes: [], scopes_supported: nil, audience: nil,
+                       jwks_uri: nil, algorithms: JWTVerifier::DEFAULT_ALGORITHMS, introspection: nil,
+                       verify_token: nil, &principal_block)
+          raise ArgumentError, "authenticate :oauth requires resource:" if resource.nil?
+          raise ArgumentError, "authenticate :oauth requires authorization_servers:" if Array(authorization_servers).empty?
+
+          validate_resource!(resource) unless resource.respond_to?(:call)
+
+          @resource = resource
+          @authorization_servers = authorization_servers
+          @scopes = Array(scopes).map(&:to_s).freeze
+          @scopes_supported = scopes_supported && Array(scopes_supported).map(&:to_s).freeze
+          @audience = audience
+          @principal_block = principal_block
+          @verifier = build_verifier(jwks_uri, algorithms, introspection, verify_token)
+        end
+
+        # Returns an Authenticator::Result. `error` is nil when no token was
+        # sent, :invalid_token or :insufficient_scope (RFC 6750 section 3.1).
+        def authenticate(request)
+          token = request.bearer_token
+          return failure(nil) if token.nil?
+
+          claims = verify(token, request)
+          return failure(:invalid_token) unless claims.is_a?(Hash)
+
+          access_token = AccessToken.new(token, claims)
+          return failure(:invalid_token) unless valid?(access_token)
+          return failure(:insufficient_scope) unless access_token.scope?(scopes)
+
+          principal = @principal_block ? @principal_block.call(access_token, request) : access_token
+          return failure(:invalid_token) unless principal
+
+          Authenticator::Result.new(true, principal, nil, access_token)
+        rescue VerificationError => e
+          MCP.logger&.error("[OmniAgent::MCP] OAuth token verification failed: #{e.message}")
+          failure(:invalid_token)
+        end
+
+        def metadata(resource_name: nil)
+          supported = @scopes_supported || scopes
+
+          {
+            resource: resource,
+            authorization_servers: authorization_servers,
+            scopes_supported: supported.empty? ? nil : supported,
+            bearer_methods_supported: [ "header" ],
+            resource_name: resource_name
+          }.compact
+        end
+
+        # Value of the WWW-Authenticate header for a 401 or 403 response.
+        def challenge(realm:, error: nil)
+          params = { realm: realm, resource_metadata: metadata_url }
+          params[:error] = error.to_s if error
+          params[:scope] = scopes.join(" ") unless scopes.empty?
+
+          "Bearer " + params.map { |key, value| %(#{key}="#{value}") }.join(", ")
+        end
+
+        # Served by RackApp under the MCP endpoint itself, so it works wherever
+        # the server is mounted.
+        def metadata_url
+          "#{resource.chomp('/')}#{METADATA_PATH}"
+        end
+
+        def resource
+          OAuth.resolve(@resource).to_s
+        end
+
+        def authorization_servers
+          Array(OAuth.resolve(@authorization_servers)).map(&:to_s)
+        end
+
+        private
+
+        def build_verifier(jwks_uri, algorithms, introspection, verify_token)
+          given = { jwks_uri: jwks_uri, introspection: introspection, verify_token: verify_token }.compact
+          unless given.size == 1
+            raise ArgumentError, "authenticate :oauth requires exactly one of jwks_uri:, introspection: or verify_token:"
+          end
+
+          if jwks_uri
+            JWTVerifier.new(jwks_uri, algorithms: algorithms, leeway: LEEWAY)
+          elsif introspection
+            IntrospectionVerifier.new(**introspection.transform_keys(&:to_sym))
+          else
+            raise ArgumentError, "verify_token: must respond to #call" unless verify_token.respond_to?(:call)
+
+            verify_token
+          end
+        end
+
+        def verify(token, request)
+          if @verifier.respond_to?(:arity) && @verifier.arity == 1
+            @verifier.call(token)
+          else
+            @verifier.call(token, request)
+          end
+        end
+
+        def valid?(access_token)
+          claims = access_token.claims
+
+          fresh?(claims) && trusted_issuer?(claims) && intended_audience?(access_token)
+        end
+
+        def fresh?(claims)
+          now = Time.now.to_i
+          return false if claims.key?("exp") && !(claims["exp"].is_a?(Numeric) && now < claims["exp"] + LEEWAY)
+          return false if claims.key?("nbf") && !(claims["nbf"].is_a?(Numeric) && now >= claims["nbf"] - LEEWAY)
+
+          true
+        end
+
+        # JWTs must carry `iss`; introspection responses and custom verifiers may omit it.
+        def trusted_issuer?(claims)
+          return true unless claims.key?("iss")
+
+          authorization_servers.map { |server| normalize(server) }.include?(normalize(claims["iss"]))
+        end
+
+        # RFC 8707: the token must have been issued for this server. Tokens
+        # meant for another resource are rejected, never passed through.
+        def intended_audience?(access_token)
+          expected = Array(OAuth.resolve(@audience) || resource).map { |value| normalize(value) }
+
+          access_token.audience.any? { |value| expected.include?(normalize(value)) }
+        end
+
+        def normalize(uri)
+          uri.to_s.chomp("/")
+        end
+
+        def validate_resource!(resource)
+          uri = URI.parse(resource.to_s)
+          return if uri.is_a?(URI::HTTP) && uri.host && uri.fragment.nil?
+
+          raise ArgumentError, "resource: must be the absolute http(s) URL of the MCP endpoint, without a fragment"
+        rescue URI::InvalidURIError
+          raise ArgumentError, "resource: #{resource.inspect} is not a valid URL"
+        end
+
+        def failure(error)
+          Authenticator::Result.new(false, nil, error, nil)
+        end
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/oauth.rb

@@ -0,0 +1,18 @@
+module OmniAgent
+  module MCP
+    # Resource server side of the MCP authorization spec (OAuth 2.1). Access
+    # tokens are issued by an external authorization server and validated here;
+    # see ResourceServer.
+    module OAuth
+      # The authorization server could not be reached or answered with something
+      # unusable. The request is rejected as unauthenticated.
+      class VerificationError < OmniAgent::MCPError; end
+
+      # Settings accept a value or a callable, so they can be read from ENV or
+      # credentials on every request.
+      def self.resolve(value)
+        value.respond_to?(:call) ? value.call : value
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/rack_app.rb

@@ -0,0 +1,131 @@
+require "json"
+
+module OmniAgent
+  module MCP
+    # Streamable HTTP transport, stateless: every POST carries one JSON-RPC
+    # message and gets a plain JSON response. No SSE stream, no sessions.
+    class RackApp
+      JSON_HEADERS = { "content-type" => "application/json" }.freeze
+      PARSE_FAILED = Object.new.freeze
+      METADATA_CORS_HEADERS = {
+        "access-control-allow-origin" => "*",
+        "access-control-allow-methods" => "GET, OPTIONS",
+        "access-control-allow-headers" => "mcp-protocol-version"
+      }.freeze
+
+      # Accepts the server class or its name. Passing a String defers the
+      # constant lookup to request time, which keeps routes reload-safe.
+      attr_reader :server
+
+      def initialize(server)
+        @server = server
+      end
+
+      def inspect
+        "#<#{self.class.name} #{server}>"
+      end
+
+      def call(env)
+        server_class = MCP.resolve_server_class(@server)
+        request = Request.new(env)
+        resource_server = server_class.configured_oauth_resource_server
+
+        # Public, so clients can discover the authorization server before they have a token.
+        if resource_server && request.path_info == OAuth::ResourceServer::METADATA_PATH
+          return protected_resource_metadata(server_class, resource_server, request)
+        end
+
+        return forbidden unless origin_allowed?(request, server_class)
+
+        auth = Authenticator.new(server_class).authenticate(request)
+        return unauthorized(server_class, resource_server, auth) unless auth.authenticated?
+
+        return method_not_allowed unless request.post?
+
+        if request.protocol_version && !SUPPORTED_PROTOCOL_VERSIONS.include?(request.protocol_version)
+          return http_error(400, "Unsupported MCP-Protocol-Version: #{request.protocol_version}")
+        end
+
+        body = read_body(request)
+        return http_error(413, "Request body too large") if body.nil?
+
+        message = parse_json(body)
+        return json_response(400, parse_error_response) if message.equal?(PARSE_FAILED)
+
+        server = server_class.new(principal: auth.principal, access_token: auth.access_token, request: request, transport: :http)
+        response = Handler.new(server).call(message)
+        return [ 202, {}, [] ] if response.nil?
+
+        status = response[:error] && response[:id].nil? ? 400 : 200
+        json_response(status, response)
+      end
+
+      private
+
+      def origin_allowed?(request, server_class)
+        origin = request.origin
+        return true if origin.nil? || origin.empty?
+
+        allowed = server_class.configured_allowed_origins + Array(OmniAgent.configuration.mcp_allowed_origins)
+        allowed.any? do |pattern|
+          pattern.is_a?(Regexp) ? pattern.match?(origin) : pattern.to_s == origin
+        end
+      end
+
+      def read_body(request)
+        max_bytes = OmniAgent.configuration.mcp_max_request_bytes
+        return nil if request.content_length.to_i > max_bytes
+
+        io = request.body
+        return "" if io.nil?
+
+        body = io.read(max_bytes + 1).to_s
+        body.bytesize > max_bytes ? nil : body
+      end
+
+      def parse_json(body)
+        JSON.parse(body)
+      rescue JSON::ParserError
+        PARSE_FAILED
+      end
+
+      def parse_error_response
+        { jsonrpc: JSONRPC_VERSION, id: nil, error: { code: PARSE_ERROR, message: "Parse error" } }
+      end
+
+      def json_response(status, payload)
+        [ status, JSON_HEADERS.dup, [ JSON.generate(payload) ] ]
+      end
+
+      def http_error(status, message)
+        json_response(status, { error: message })
+      end
+
+      def forbidden
+        http_error(403, "Origin not allowed")
+      end
+
+      def method_not_allowed(allow = "POST")
+        [ 405, JSON_HEADERS.merge("allow" => allow), [ JSON.generate({ error: "Method not allowed" }) ] ]
+      end
+
+      def unauthorized(server_class, resource_server, auth)
+        realm = server_class.server_name
+        challenge = resource_server ? resource_server.challenge(realm: realm, error: auth.error) : %(Bearer realm="#{realm}")
+        status, message = auth.error == :insufficient_scope ? [ 403, "Insufficient scope" ] : [ 401, "Unauthorized" ]
+
+        [ status, JSON_HEADERS.merge("www-authenticate" => challenge), [ JSON.generate({ error: message }) ] ]
+      end
+
+      # RFC 9728 Protected Resource Metadata. Browser-based clients fetch it
+      # cross-origin, and it holds nothing secret, so any origin may read it.
+      def protected_resource_metadata(server_class, resource_server, request)
+        return [ 204, METADATA_CORS_HEADERS.dup, [] ] if request.options?
+        return method_not_allowed("GET") unless request.get? || request.head?
+
+        metadata = resource_server.metadata(resource_name: server_class.server_name)
+        [ 200, JSON_HEADERS.merge(METADATA_CORS_HEADERS), [ JSON.generate(metadata) ] ]
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/request.rb

@@ -0,0 +1,24 @@
+require "rack"
+
+module OmniAgent
+  module MCP
+    class Request < Rack::Request
+      def bearer_token
+        header = get_header("HTTP_AUTHORIZATION").to_s
+        match = header.match(/\ABearer\s+(.+)\z/i)
+        return if match.nil?
+
+        token = match[1].strip
+        token.empty? ? nil : token
+      end
+
+      def protocol_version
+        get_header("HTTP_MCP_PROTOCOL_VERSION")
+      end
+
+      def origin
+        get_header("HTTP_ORIGIN")
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/routing.rb

@@ -0,0 +1,57 @@
+require "active_support/core_ext/string/inflections"
+
+module OmniAgent
+  module MCP
+    # Router helpers, included into ActionDispatch::Routing::Mapper by the engine.
+    #
+    #   mcp_server :support                        # SupportServer at /mcp/support
+    #   mcp_server :support, path: "/internal/mcp" # custom path
+    #   mcp_servers :support, :billing             # several at once
+    #   mcp_servers                                # every app/mcp_servers/**/*_server.rb
+    module Routing
+      DEFAULT_PATH_PREFIX = "/mcp".freeze
+
+      def mcp_server(name, path: nil, server: nil, as: nil)
+        slug = OmniAgent::MCP::Routing.slug_for(name)
+        server_name = server.is_a?(Class) ? server.name : (server || "#{slug}_server".camelize).to_s
+
+        mount OmniAgent::MCP::RackApp.new(server_name) => (path || "#{DEFAULT_PATH_PREFIX}/#{slug.dasherize}"),
+              as: (as || "#{slug.tr('/', '_')}_mcp_server")
+      end
+
+      def mcp_servers(*names, except: [], directory: nil)
+        names = names.flatten
+        names = OmniAgent::MCP::Routing.discover(directory || Rails.root.join("app", "mcp_servers")) if names.empty?
+        excluded = Array(except).map { |name| OmniAgent::MCP::Routing.slug_for(name) }
+
+        names.each do |name|
+          mcp_server(name) unless excluded.include?(OmniAgent::MCP::Routing.slug_for(name))
+        end
+      end
+
+      # :support, "support", "support_server", "SupportServer" -> "support"
+      # "Admin::SupportServer", "admin/support" -> "admin/support"
+      def self.slug_for(name)
+        name.to_s.underscore.delete_suffix("_server")
+      end
+
+      # Server slugs for every *_server.rb file, without loading any code.
+      # Skips files inside a server's own directory (e.g. support_server/tools/)
+      # and abstract base servers named application_*.
+      def self.discover(directory)
+        root = Pathname.new(directory.to_s)
+        return [] unless root.directory?
+
+        Dir.glob(root.join("**", "*_server.rb").to_s).sort.filter_map do |file|
+          relative = Pathname.new(file).relative_path_from(root).to_s.delete_suffix(".rb")
+          *dirs, base = relative.split("/")
+
+          next if dirs.any? { |dir| dir.end_with?("_server") }
+          next if base.start_with?("application_")
+
+          slug_for(relative)
+        end
+      end
+    end
+  end
+end