omni_agent 0.1.9 → 0.1.11

data/lib/omni_agent/mcp/server.rb

@@ -0,0 +1,255 @@
+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, **oauth_options, &block)
+          if strategy != :oauth && oauth_options.any?
+            raise ArgumentError, "unknown keywords for authenticate: #{oauth_options.keys.map(&:inspect).join(', ')}"
+          end
+
+          @authentication = if strategy == :oauth
+            raise ArgumentError, "authenticate :oauth does not take tokens:" unless tokens.nil?
+
+            # The block, if any, maps the validated access token to the principal.
+            { strategy: :oauth, resource_server: OAuth::ResourceServer.new(**oauth_options, &block) }
+          elsif 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:), :oauth, :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
+
+        def configured_oauth_resource_server
+          authentication = configured_authentication
+          authentication[:resource_server] if authentication && authentication[:strategy] == :oauth
+        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
+
+      # access_token is the OAuth::AccessToken for `authenticate :oauth`, nil otherwise.
+      attr_reader :principal, :request, :transport, :access_token
+
+      def initialize(principal: nil, request: nil, transport: :http, access_token: nil)
+        @principal = principal
+        @request = request
+        @transport = transport
+        @access_token = access_token
+      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 }
+        base[:mcp_access_token] = access_token if access_token
+        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/tool.rb

@@ -11,6 +11,31 @@ module OmniAgent
         @metadata || {}
       end
 
+      def title(text = nil)
+        @title = text if text
+        @title
+      end
+
+      def annotations(read_only: nil, destructive: nil, idempotent: nil, open_world: nil)
+        @annotations = {
+          read_only: read_only,
+          destructive: destructive,
+          idempotent: idempotent,
+          open_world: open_world
+        }.compact
+      end
+
+      def mcp_annotations
+        hints = @annotations || {}
+
+        {
+          readOnlyHint: hints[:read_only],
+          destructiveHint: hints[:destructive],
+          idempotentHint: hints[:idempotent],
+          openWorldHint: hints[:open_world]
+        }.compact
+      end
+
       def tags(*tag_names)
         return @tags || [] if tag_names.empty?
 

data/lib/omni_agent/version.rb

@@ -1,3 +1,3 @@
 module OmniAgent
-  VERSION = "0.1.9"
+  VERSION = "0.1.11"
 end

data/lib/omni_agent.rb

@@ -3,7 +3,10 @@ require "omni_agent/engine" if defined?(Rails)
 require "zeitwerk"
 
 loader = Zeitwerk::Loader.for_gem
-loader.inflector.inflect("openai" => "OpenAI", "cli" => "CLI")
+loader.inflector.inflect(
+  "openai" => "OpenAI", "cli" => "CLI", "mcp" => "MCP",
+  "oauth" => "OAuth", "jwt_verifier" => "JWTVerifier", "http_client" => "HTTPClient"
+)
 loader.ignore(File.expand_path("generators", __dir__))
 loader.setup
 

data/lib/tasks/omni_agent_mcp_tasks.rake

@@ -0,0 +1,9 @@
+namespace :omni_agent do
+  desc 'Run an MCP server over stdio. Usage: rake "omni_agent:mcp[SupportServer]"'
+  task :mcp, [ :server ] => :environment do |_task, args|
+    server = args[:server].to_s.strip
+    abort 'Please provide a server class. Example: bin/rails "omni_agent:mcp[SupportServer]"' if server.empty?
+
+    OmniAgent::MCP::Stdio.start(server)
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omni_agent
 version: !ruby/object:Gem::Version
-  version: 0.1.9
+  version: 0.1.11
 platform: ruby
 authors:
 - ACR1209
@@ -62,6 +62,7 @@ files:
 - lib/generators/omni_agent/agent/agent_generator.rb
 - lib/generators/omni_agent/eval/eval_generator.rb
 - lib/generators/omni_agent/install/install_generator.rb
+- lib/generators/omni_agent/mcp_server/mcp_server_generator.rb
 - lib/omni_agent.rb
 - lib/omni_agent/agent.rb
 - lib/omni_agent/configuration.rb
@@ -81,6 +82,21 @@ files:
 - lib/omni_agent/eval/run.rb
 - lib/omni_agent/eval/runner.rb
 - lib/omni_agent/eval/tool_call_assertion.rb
+- lib/omni_agent/mcp.rb
+- lib/omni_agent/mcp/authenticator.rb
+- lib/omni_agent/mcp/handler.rb
+- lib/omni_agent/mcp/oauth.rb
+- lib/omni_agent/mcp/oauth/access_token.rb
+- lib/omni_agent/mcp/oauth/http_client.rb
+- lib/omni_agent/mcp/oauth/introspection_verifier.rb
+- lib/omni_agent/mcp/oauth/jwt_verifier.rb
+- lib/omni_agent/mcp/oauth/resource_server.rb
+- lib/omni_agent/mcp/rack_app.rb
+- lib/omni_agent/mcp/request.rb
+- lib/omni_agent/mcp/routing.rb
+- lib/omni_agent/mcp/server.rb
+- lib/omni_agent/mcp/stdio.rb
+- lib/omni_agent/mcp/tool_adapter.rb
 - lib/omni_agent/providers.rb
 - lib/omni_agent/providers/base.rb
 - lib/omni_agent/providers/mock.rb
@@ -96,6 +112,7 @@ files:
 - lib/omni_agent/tool/schema_builder.rb
 - lib/omni_agent/version.rb
 - lib/tasks/omni_agent_eval_tasks.rake
+- lib/tasks/omni_agent_mcp_tasks.rake
 homepage: https://github.com/ACR1209/omni_agent
 licenses:
 - MIT