omni_agent 0.1.8 → 0.1.10

data/lib/omni_agent/mcp/handler.rb

@@ -0,0 +1,128 @@
+module OmniAgent
+  module MCP
+    class Handler
+      class ProtocolError < StandardError
+        attr_reader :code
+
+        def initialize(code, message)
+          super(message)
+          @code = code
+        end
+      end
+
+      def initialize(server)
+        @server = server
+      end
+
+      # Takes a parsed JSON-RPC message and returns the response Hash, or nil
+      # when no response is due (notifications and client responses).
+      def call(message)
+        unless message.is_a?(Hash)
+          reason = message.is_a?(Array) ? "Batch requests are not supported" : "Invalid Request"
+          return error_response(nil, INVALID_REQUEST, reason)
+        end
+
+        message = message.transform_keys(&:to_s)
+        return nil if client_response?(message)
+
+        id = message["id"]
+        notification = !message.key?("id")
+
+        unless message["jsonrpc"] == JSONRPC_VERSION && message["method"].is_a?(String)
+          return notification ? nil : error_response(id, INVALID_REQUEST, "Invalid Request")
+        end
+
+        return nil if notification
+
+        result = dispatch(message["method"], message["params"])
+        success_response(id, result)
+      rescue ProtocolError => e
+        error_response(id, e.code, e.message)
+      rescue StandardError => e
+        log_error("MCP internal error", e)
+        error_response(id, INTERNAL_ERROR, "Internal error")
+      end
+
+      private
+
+      attr_reader :server
+
+      def dispatch(method, params)
+        params = normalize_params(params)
+
+        case method
+        when "initialize" then initialize_result(params)
+        when "ping" then {}
+        when "tools/list" then tools_list_result
+        when "tools/call" then tools_call_result(params)
+        else
+          raise ProtocolError.new(METHOD_NOT_FOUND, "Method not found: #{method}")
+        end
+      end
+
+      def normalize_params(params)
+        return {} if params.nil?
+        raise ProtocolError.new(INVALID_PARAMS, "params must be an object") unless params.is_a?(Hash)
+
+        params.transform_keys(&:to_s)
+      end
+
+      def initialize_result(params)
+        requested = params["protocolVersion"]
+        protocol_version = SUPPORTED_PROTOCOL_VERSIONS.include?(requested) ? requested : LATEST_PROTOCOL_VERSION
+        server_class = server.class
+
+        result = {
+          protocolVersion: protocol_version,
+          capabilities: { tools: { listChanged: false } },
+          serverInfo: { name: server_class.server_name, version: server_class.server_version }
+        }
+        result[:instructions] = server_class.instructions if server_class.instructions
+        result
+      end
+
+      def tools_list_result
+        tools = server.visible_tools.map { |name, tool_class| ToolAdapter.definition(name, tool_class) }
+        { tools: tools }
+      end
+
+      def tools_call_result(params)
+        name = params["name"]
+        raise ProtocolError.new(INVALID_PARAMS, "Missing tool name") unless name.is_a?(String)
+
+        tool_class = server.find_tool(name)
+        raise ProtocolError.new(INVALID_PARAMS, "Unknown tool: #{name}") unless tool_class
+
+        arguments = params["arguments"] || {}
+        raise ProtocolError.new(INVALID_PARAMS, "arguments must be an object") unless arguments.is_a?(Hash)
+
+        execute_tool(name, tool_class, arguments)
+      end
+
+      def execute_tool(name, tool_class, arguments)
+        tool = tool_class.new
+        tool.context = server.tool_context
+        ToolAdapter.result(tool.invoke(arguments))
+      rescue StandardError => e
+        log_error("MCP tool #{name} failed", e)
+        ToolAdapter.error_result("Error executing tool: #{e.message}")
+      end
+
+      def client_response?(message)
+        !message.key?("method") && (message.key?("result") || message.key?("error"))
+      end
+
+      def success_response(id, result)
+        { jsonrpc: JSONRPC_VERSION, id: id, result: result }
+      end
+
+      def error_response(id, code, message)
+        { jsonrpc: JSONRPC_VERSION, id: id, error: { code: code, message: message } }
+      end
+
+      def log_error(prefix, error)
+        MCP.logger&.error("[OmniAgent::MCP] #{prefix}: #{error.class}: #{error.message}")
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/rack_app.rb

@@ -0,0 +1,107 @@
+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
+
+      # 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)
+
+        return forbidden unless origin_allowed?(request, server_class)
+
+        auth = Authenticator.new(server_class).authenticate(request)
+        return unauthorized(server_class) 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, 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
+        [ 405, JSON_HEADERS.merge("allow" => "POST"), [ JSON.generate({ error: "Method not allowed" }) ] ]
+      end
+
+      def unauthorized(server_class)
+        headers = JSON_HEADERS.merge("www-authenticate" => %(Bearer realm="#{server_class.server_name}"))
+        [ 401, headers, [ JSON.generate({ error: "Unauthorized" }) ] ]
+      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

data/lib/omni_agent/mcp/server.rb

@@ -0,0 +1,238 @@
+require "active_support/core_ext/string/inflections"
+
+module OmniAgent
+  module MCP
+    class Server
+      TOOL_NAME_PATTERN = /\A[A-Za-z0-9_.-]{1,128}\z/
+
+      class << self
+        def server_name(value = nil)
+          @server_name = value.to_s if value
+          setting(:@server_name) || default_server_name
+        end
+
+        def server_version(value = nil)
+          @server_version = value.to_s if value
+          setting(:@server_version) || "0.1.0"
+        end
+
+        def instructions(text = nil)
+          @instructions = text if text
+          setting(:@instructions)
+        end
+
+        def tools(*tool_classes)
+          tool_classes.flatten.each { |tool_class| tool(tool_class) }
+        end
+
+        def tool(tool_class, as: nil)
+          add_tool_declaration(type: :tool, tool_class: tool_class, name: as&.to_s)
+        end
+
+        def tools_from(*agent_classes)
+          agent_classes.flatten.each do |agent_class|
+            add_tool_declaration(type: :agent_tools, agent_class: agent_class)
+          end
+        end
+
+        def expose_agent(agent_class, as:, description: nil, run_alias: nil, forward: [])
+          unless agent_class.is_a?(Class) && agent_class <= OmniAgent::Agent
+            raise ArgumentError, "expose_agent requires an OmniAgent::Agent subclass"
+          end
+
+          tool_class = OmniAgent::Agent.build_delegated_tool_class(
+            agent_class,
+            description: description,
+            run_alias: run_alias,
+            forward: forward
+          )
+
+          add_tool_declaration(type: :tool, tool_class: tool_class, name: as.to_s)
+        end
+
+        def authenticate(strategy = nil, tokens: nil, &block)
+          @authentication = if block
+            { strategy: :custom, block: block }
+          elsif strategy == :none
+            { strategy: :none }
+          elsif strategy == :bearer
+            raise ArgumentError, "authenticate :bearer requires tokens:" if tokens.nil?
+
+            { strategy: :bearer, tokens: tokens }
+          else
+            raise ArgumentError, "authenticate expects :bearer (with tokens:), :none, or a block"
+          end
+        end
+
+        def authorize_tool(&block)
+          raise ArgumentError, "authorize_tool requires a block" unless block
+
+          @authorize_tool_block = block
+        end
+
+        def context(&block)
+          raise ArgumentError, "context requires a block" unless block
+
+          @context_block = block
+        end
+
+        def stdio_principal(&block)
+          raise ArgumentError, "stdio_principal requires a block" unless block
+
+          @stdio_principal_block = block
+        end
+
+        def allowed_origins(*origins)
+          @allowed_origins = (@allowed_origins || []) + origins.flatten
+        end
+
+        def configured_tool_declarations; list_setting(:@tool_declarations); end
+        def configured_authentication; setting(:@authentication); end
+        def configured_authorize_tool_block; setting(:@authorize_tool_block); end
+        def configured_context_block; setting(:@context_block); end
+        def configured_stdio_principal_block; setting(:@stdio_principal_block); end
+        def configured_allowed_origins; list_setting(:@allowed_origins); end
+
+        # Ordered { "ToolName" => tool_class }: for each server in the
+        # inheritance chain, its <Server>::Tools namespace first, then its
+        # declared tools. Resolved on every call so lazily referenced tools
+        # (strings, tools_from, the Tools namespace) pick up code reloads.
+        def tool_registry
+          server_ancestors.reverse.each_with_object({}) do |server_class, registry|
+            entries = server_class.namespace_tool_classes.map { |tool_class| [ default_tool_name(tool_class), tool_class ] }
+            entries += (server_class.instance_variable_get(:@tool_declarations) || []).flat_map do |declaration|
+              resolve_declaration(declaration)
+            end
+
+            entries.each { |name, tool_class| register_tool(registry, name, tool_class) }
+          end
+        end
+
+        # Tools defined under this server's own namespace, e.g.
+        # app/mcp_servers/support_server/tools/*.rb -> SupportServer::Tools::*.
+        def namespace_tool_classes
+          return [] unless const_defined?(:Tools, false)
+
+          namespace = const_get(:Tools, false)
+          return [] unless namespace.is_a?(Module)
+
+          namespace.constants.sort.filter_map do |const_name|
+            const = namespace.const_get(const_name, false)
+            const if const.is_a?(Class) && const < OmniAgent::Tool
+          end
+        end
+
+        def resolve_stdio_principal
+          configured_stdio_principal_block&.call
+        end
+
+        private
+
+        # Servers inherit settings from parent servers (e.g. an ApplicationMCPServer
+        # that declares `authenticate` once). Scalars: nearest wins. Lists: concatenated.
+        def server_ancestors
+          ancestors.select { |klass| klass.is_a?(Class) && klass < OmniAgent::MCP::Server }
+        end
+
+        def setting(ivar)
+          server_ancestors.each do |klass|
+            value = klass.instance_variable_get(ivar)
+            return value unless value.nil?
+          end
+
+          nil
+        end
+
+        def list_setting(ivar)
+          server_ancestors.reverse.flat_map { |klass| klass.instance_variable_get(ivar) || [] }
+        end
+
+        def default_server_name
+          return "omni_agent" if name.nil?
+
+          base = name.demodulize.underscore
+          base = base.delete_suffix("_server") unless base == "server"
+          base
+        end
+
+        def add_tool_declaration(declaration)
+          @tool_declarations = (@tool_declarations || []) + [ declaration ]
+        end
+
+        def resolve_declaration(declaration)
+          case declaration[:type]
+          when :tool
+            tool_class = resolve_constant(declaration[:tool_class])
+            [ [ declaration[:name] || default_tool_name(tool_class), tool_class ] ]
+          when :agent_tools
+            agent_class = resolve_constant(declaration[:agent_class])
+            agent_class.tool_classes.map { |tool_class| [ default_tool_name(tool_class), tool_class ] }
+          end
+        end
+
+        def resolve_constant(value)
+          return value if value.is_a?(Class)
+
+          Object.const_get(value.to_s)
+        end
+
+        def default_tool_name(tool_class)
+          if tool_class.name.nil?
+            raise OmniAgent::MCPError, "Anonymous tool classes need an explicit name: `tool klass, as: \"name\"`"
+          end
+
+          tool_class.name.demodulize
+        end
+
+        def register_tool(registry, name, tool_class)
+          unless tool_class.is_a?(Class) && tool_class < OmniAgent::Tool
+            raise OmniAgent::MCPError, "#{tool_class.inspect} is not an OmniAgent::Tool subclass"
+          end
+
+          unless name.match?(TOOL_NAME_PATTERN)
+            raise OmniAgent::MCPError,
+                  "Invalid MCP tool name #{name.inspect} in #{self.name}: use 1-128 characters from A-Z, a-z, 0-9, _, -, ."
+          end
+
+          if registry.key?(name)
+            raise OmniAgent::MCPError,
+                  "Duplicate MCP tool name #{name.inspect} in #{self.name}. Use `tool klass, as: \"other_name\"` to rename one."
+          end
+
+          registry[name] = tool_class
+        end
+      end
+
+      attr_reader :principal, :request, :transport
+
+      def initialize(principal: nil, request: nil, transport: :http)
+        @principal = principal
+        @request = request
+        @transport = transport
+      end
+
+      def visible_tools
+        authorizer = self.class.configured_authorize_tool_block
+        registry = self.class.tool_registry
+        return registry unless authorizer
+
+        registry.select do |name, tool_class|
+          instance_exec(tool_class, principal, name, &authorizer)
+        end
+      end
+
+      def find_tool(name)
+        visible_tools[name.to_s]
+      end
+
+      def tool_context
+        base = { mcp_principal: principal, mcp_request: request }
+        context_block = self.class.configured_context_block
+        return base unless context_block
+
+        extra = instance_exec(principal, request, &context_block)
+        extra.is_a?(Hash) ? base.merge(extra) : base
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/stdio.rb

@@ -0,0 +1,56 @@
+require "json"
+
+module OmniAgent
+  module MCP
+    # stdio transport: newline-delimited JSON-RPC over stdin/stdout. Runs as a
+    # local subprocess of the client, so it skips HTTP authentication and takes
+    # its principal from the server's `stdio_principal` block instead.
+    class Stdio
+      # Entry point for `omni_agent mcp` and `rake omni_agent:mcp`. Keeps a
+      # private handle on the real stdout for the protocol, then points fd 1
+      # at stderr so stray `puts` or logger output cannot corrupt the stream.
+      def self.start(server, input: $stdin)
+        server_class = MCP.resolve_server_class(server)
+
+        protocol_output = $stdout.dup
+        $stdout.flush
+        $stdout.reopen($stderr)
+
+        new(server_class, input: input, output: protocol_output).run
+      end
+
+      def initialize(server, input:, output:)
+        @server_class = MCP.resolve_server_class(server)
+        @input = input
+        @output = output
+      end
+
+      def run
+        server = @server_class.new(principal: @server_class.resolve_stdio_principal, transport: :stdio)
+        handler = Handler.new(server)
+
+        @input.each_line do |line|
+          line = line.strip
+          next if line.empty?
+
+          message = begin
+            JSON.parse(line)
+          rescue JSON::ParserError
+            write(jsonrpc: JSONRPC_VERSION, id: nil, error: { code: PARSE_ERROR, message: "Parse error" })
+            next
+          end
+
+          response = handler.call(message)
+          write(response) if response
+        end
+      end
+
+      private
+
+      def write(payload)
+        @output.write("#{JSON.generate(payload)}\n")
+        @output.flush
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp/tool_adapter.rb

@@ -0,0 +1,45 @@
+require "json"
+
+module OmniAgent
+  module MCP
+    module ToolAdapter
+      module_function
+
+      def definition(name, tool_class)
+        definition = { name: name }
+        definition[:title] = tool_class.title if tool_class.title
+        definition[:description] = tool_class.description
+        definition[:inputSchema] = tool_class.json_schema
+
+        annotations = tool_class.mcp_annotations
+        definition[:annotations] = annotations unless annotations.empty?
+
+        definition
+      end
+
+      def result(value)
+        case value
+        when nil
+          { content: [ text_content("") ] }
+        when String
+          { content: [ text_content(value) ] }
+        when Hash
+          json = JSON.generate(value)
+          { content: [ text_content(json) ], structuredContent: JSON.parse(json) }
+        when Array
+          { content: [ text_content(JSON.generate(value)) ] }
+        else
+          { content: [ text_content(value.to_s) ] }
+        end
+      end
+
+      def error_result(message)
+        { content: [ text_content(message) ], isError: true }
+      end
+
+      def text_content(text)
+        { type: "text", text: text }
+      end
+    end
+  end
+end

data/lib/omni_agent/mcp.rb

@@ -0,0 +1,31 @@
+module OmniAgent
+  module MCP
+    SUPPORTED_PROTOCOL_VERSIONS = [ "2025-11-25", "2025-06-18", "2025-03-26" ].freeze
+    LATEST_PROTOCOL_VERSION = SUPPORTED_PROTOCOL_VERSIONS.first
+
+    JSONRPC_VERSION = "2.0".freeze
+
+    PARSE_ERROR = -32_700
+    INVALID_REQUEST = -32_600
+    METHOD_NOT_FOUND = -32_601
+    INVALID_PARAMS = -32_602
+    INTERNAL_ERROR = -32_603
+
+    def self.logger
+      return unless defined?(Rails) && Rails.respond_to?(:logger)
+
+      Rails.logger
+    end
+
+    def self.resolve_server_class(server)
+      return server if server.is_a?(Class)
+
+      server_class = Object.const_get(server.to_s)
+      unless server_class.is_a?(Class) && server_class < OmniAgent::MCP::Server
+        raise OmniAgent::MCPError, "#{server} is not an OmniAgent::MCP::Server subclass"
+      end
+
+      server_class
+    end
+  end
+end

data/lib/omni_agent/providers/base.rb

@@ -9,7 +9,7 @@ module OmniAgent
         @model = model || default_model
       end
 
-      def chat(messages:, tools: [], **_options)
+      def chat(messages:, tools: [], stream: nil, **_options)
         raise NotImplementedError, "Providers must implement #chat"
       end