toolchest 0.3.6 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c73baf6206cfd9d1f744867f2b3e35e6094a2f4304dec858ce0350c1d54c03ba
-  data.tar.gz: 82c480aa62bdcb7c19df80a1aebeb4a7cf4dbb4ec8a199b25fb4a8abe6da69a1
+  metadata.gz: 80dd8a25c340a7918653dc8e63fa0d899a23e10306c87d3143456688891de9d7
+  data.tar.gz: 0f2416b46244ff3e03aa28636e81d126e18303f8385bea0053b4e90c04c44551
 SHA512:
-  metadata.gz: 82bcf0c2b9ee5c5d0e9365641dcd7e2e6119be86e66dadee37b1bfaa44582511b2c3e7a5f2cbf31b0164dd31d4c3823090900bb4dd9bf18ea9f215babeaa2e47
-  data.tar.gz: f0e79e3d5cd3ad4bd09fe68975e129a3528f78f5f5ec7eaa822bf37c85631a295df2d0e10acbaf483ee62fa243ce563bd1a264688e3f3b90e1884a96d5173e1c
+  metadata.gz: ae68328956e828a05ae03e662d95f8af5052de7d8a456a06d99c76a2989665f67b236003ef990b9491cec6e7b85310dabfd8e5bdc3a35584189c0710ebf48c71
+  data.tar.gz: 446511d71d09c9d1c02728176d559d478ff15df18ce317b118bd1d1e6fe3f9e291cc10ba71c7a1fc626493c991bbdb5c57f74321f288715bad98384e562cbe43

data/README.md

@@ -278,6 +278,79 @@ tool "Export data", access: :read, annotations: { openWorldHint: true } do
 end
 ```
 
+### Elicitation
+
+Ask the client's user for input via a form. Returns the result hash from the client.
+
+```ruby
+tool "Confirm shipping address" do
+  param :order_id, :string, "Order ID"
+end
+def confirm_address
+  @order = Order.find(params[:order_id])
+  result = mcp_elicit("Please confirm the shipping address: #{@order.address}",
+    schema: {
+      type: "object",
+      properties: {
+        confirmed: { type: "boolean", description: "Address is correct?" }
+      },
+      required: ["confirmed"]
+    })
+  halt error: "Cancelled" unless result["action"] == "accept"
+  @order.confirm_address!
+  render :show
+end
+```
+
+Raises `Toolchest::Error` if the client doesn't support elicitation.
+
+### Cancellation
+
+Check if the client cancelled the current request. Useful in long-running loops:
+
+```ruby
+tool "Import customers" do
+  param :file_url, :string, "CSV URL"
+end
+def import
+  rows = CSV.parse(download(params[:file_url]))
+  rows.each_with_index do |row, i|
+    mcp_raise_if_cancelled!
+    Customer.create!(row.to_h)
+    mcp_progress i + 1, total: rows.size
+  end
+  render text: "Imported #{rows.size} customers"
+end
+```
+
+`mcp_cancelled?` returns a boolean. `mcp_raise_if_cancelled!` raises `MCP::CancelledError` to bail out immediately.
+
+### Tool titles and output schemas
+
+`title:` sets a human-readable display name. `output:` declares the tool's output JSON schema:
+
+```ruby
+tool "Look up an order", title: "Order Lookup",
+  output: { type: "object", properties: { id: { type: "string" }, status: { type: "string" } } } do
+  param :order_id, :string, "Order ID"
+end
+def show
+  # ...
+end
+```
+
+`title:` also works on `resource` and `prompt`:
+
+```ruby
+resource "orders://schema", name: "Order schema", title: "Schema: Orders" do
+  # ...
+end
+
+prompt "debug-order", title: "Debug Order" do |order_id:|
+  # ...
+end
+```
+
 ### Logging
 
 ```ruby
@@ -300,6 +373,51 @@ end
 
 This shows up in the MCP initialize response. `server_name` and `server_description` are also available.
 
+### Transport security
+
+Configure host validation, origin checking, and session limits:
+
+```ruby
+Toolchest.configure do |config|
+  config.allowed_hosts = ["myapp.com", "mcp.myapp.com"]
+  config.allowed_origins = ["https://myapp.com"]
+  config.dns_rebinding_protection = true
+  config.session_idle_timeout = 30.minutes
+  config.max_sessions = 100
+  config.stateless = false
+end
+```
+
+`dns_rebinding_protection` validates the `Host` header against `allowed_hosts` to prevent DNS rebinding attacks. `stateless` disables server-side session state entirely (every request is standalone). `session_idle_timeout` evicts idle sessions, `max_sessions` caps total concurrent sessions.
+
+### Caching and pagination
+
+Control list pagination and cache behavior for resources:
+
+```ruby
+Toolchest.configure do |config|
+  config.page_size = 50
+  config.cache_ttl = 5.minutes
+  config.cache_scope = "public"
+end
+```
+
+`page_size` sets the default page size for paginated list responses. `cache_ttl` and `cache_scope` control how long clients can cache list and read results. `cache_scope` is `"public"` (shared caches allowed) or `"private"` (client-only).
+
+### Instrumentation
+
+Every tool dispatch fires an `ActiveSupport::Notifications` event:
+
+```ruby
+ActiveSupport::Notifications.subscribe("dispatch.toolchest") do |*args|
+  event = ActiveSupport::Notifications::Event.new(*args)
+  Rails.logger.info "[toolchest] #{event.payload[:tool]} " \
+    "#{event.payload[:duration]}ms error=#{event.payload[:error]}"
+end
+```
+
+Payload includes `:tool`, `:toolbox`, `:action`, `:arguments`, `:duration` (ms), and `:error` (boolean).
+
 ## Auth
 
 Three built-in strategies, or bring your own. Default is `:none`.
@@ -630,10 +748,11 @@ end
 
 - **Rate limiting**: Toolchest doesn't include rate limiting. Use [rack-attack](https://github.com/rack/rack-attack) or your reverse proxy to protect token and registration endpoints.
 - **HTTPS**: OAuth endpoints should always run behind TLS in production.
+- **DNS rebinding**: Built-in protection via `config.dns_rebinding_protection` and `config.allowed_hosts`. See [Transport security](#transport-security).
 
 ## Internals
 
-Transport is the [MCP Ruby SDK](https://github.com/modelcontextprotocol/ruby-sdk) (`mcp` gem).
+Transport is the [MCP Ruby SDK](https://github.com/modelcontextprotocol/ruby-sdk) (`mcp` gem, >= 1.0).
 
 OAuth provider is cribbed from [Doorkeeper](https://github.com/doorkeeper-gem/doorkeeper). Same table layout, same controller shapes. Not a dependency, just stole the design.
 

data/app/controllers/toolchest/application_controller.rb

@@ -1,5 +1,14 @@
 module Toolchest
   class ApplicationController < Toolchest.base_controller.constantize
+    protect_from_forgery with: :exception
+    before_action :set_frame_options
+
     helper Toolchest::RouteDelegation if Toolchest.delegate_route_helpers
+
+    private
+
+    def set_frame_options
+      response.headers["X-Frame-Options"] = "DENY"
+    end
   end
 end

data/app/controllers/toolchest/oauth/authorizations_controller.rb

@@ -133,6 +133,25 @@ module Toolchest
 
         if params[:redirect_uri].present? && !@application.redirect_uri_matches?(params[:redirect_uri])
           render json: { error: "invalid_redirect_uri" }, status: :bad_request
+          return
+        end
+
+        if params[:response_type].present? && params[:response_type] != "code"
+          redirect_url = build_redirect(params[:redirect_uri],
+            error: "unsupported_response_type",
+            state: params[:state]
+          )
+          redirect_to redirect_url, allow_other_host: true
+          return
+        end
+
+        if params[:code_challenge_method].present? && params[:code_challenge_method] != "S256"
+          redirect_url = build_redirect(params[:redirect_uri],
+            error: "invalid_request",
+            error_description: "Only S256 code_challenge_method is supported",
+            state: params[:state]
+          )
+          redirect_to redirect_url, allow_other_host: true
         end
       end
 

data/app/controllers/toolchest/oauth/metadata_controller.rb

@@ -37,14 +37,9 @@ module Toolchest
       # Returns [mount_path, config] or renders 404 and returns [nil, nil].
       def resolve_mount
         if params[:rest].present?
-          # Suffixed path (RFC 8414) — must match a configured mount exactly.
           path = "/#{params[:rest]}"
           key = Toolchest.mount_keys.find { |k| Toolchest.configuration(k).mount_path == path }
-          unless key
-            head :not_found
-            return [nil, nil]
-          end
-          return [path, Toolchest.configuration(key)]
+          return [path, Toolchest.configuration(key)] if key
         end
 
         # No suffix (e.g. Cursor). Use default_oauth_mount if set.

data/app/controllers/toolchest/oauth/registrations_controller.rb

@@ -31,6 +31,30 @@ module Toolchest
           }, status: :bad_request
         end
 
+        dangerous_schemes = %w[javascript data vbscript].freeze
+        uris.each do |u|
+          begin
+            parsed = URI.parse(u)
+            unless parsed.scheme && parsed.host
+              return render json: {
+                error: "invalid_client_metadata",
+                error_description: "Redirect URI must have a scheme and host"
+              }, status: :bad_request
+            end
+            if dangerous_schemes.include?(parsed.scheme.downcase)
+              return render json: {
+                error: "invalid_client_metadata",
+                error_description: "Redirect URI scheme is not allowed"
+              }, status: :bad_request
+            end
+          rescue URI::InvalidURIError
+            return render json: {
+              error: "invalid_client_metadata",
+              error_description: "Redirect URI is not a valid URI"
+            }, status: :bad_request
+          end
+        end
+
         application = Toolchest::OauthApplication.new(
           name: name,
           redirect_uri: uris.join("\n"),

data/app/controllers/toolchest/oauth/tokens_controller.rb

@@ -21,7 +21,7 @@ module Toolchest
       def toolchest_config = Toolchest.configuration(mount_key.to_sym)
 
       def handle_authorization_code
-        grant = Toolchest::OauthAccessGrant.find_by_code(params[:code])
+        grant = Toolchest::OauthAccessGrant.find_by_code(params[:code], mount_key: mount_key)
 
         unless grant
           return error_response("invalid_grant", "Authorization code not found or expired")
@@ -29,7 +29,11 @@ module Toolchest
 
         app = grant.application
 
-        if params[:client_id].present? && app.uid != params[:client_id]
+        unless params[:client_id].present?
+          return error_response("invalid_request", "client_id is required")
+        end
+
+        if app.uid != params[:client_id]
           return error_response("invalid_client", "Client ID mismatch")
         end
 
@@ -37,8 +41,8 @@ module Toolchest
           return error_response("invalid_grant", "Redirect URI mismatch")
         end
 
-        if !grant.uses_pkce? && !app.confidential?
-          return error_response("invalid_request", "PKCE required for public clients")
+        unless grant.uses_pkce?
+          return error_response("invalid_request", "PKCE is required")
         end
 
         unless grant.verify_pkce(params[:code_verifier])
@@ -69,7 +73,13 @@ module Toolchest
           return error_response("invalid_grant", "Refresh token invalid or expired")
         end
 
-        old_token.revoke!
+        if params[:client_id].present? && old_token.application&.uid != params[:client_id]
+          return error_response("invalid_client", "Client ID mismatch")
+        end
+
+        unless old_token.revoke_atomically!
+          return error_response("invalid_grant", "Refresh token already used")
+        end
 
         token = Toolchest::OauthAccessToken.create_for(
           application: old_token.application,

data/app/models/toolchest/oauth_access_grant.rb

@@ -65,7 +65,11 @@ module Toolchest
         grant
       end
 
-      def find_by_code(raw_code) = active.find_by(token_digest: Digest::SHA256.hexdigest(raw_code))
+      def find_by_code(raw_code, mount_key: nil)
+        scope = active.where(token_digest: Digest::SHA256.hexdigest(raw_code))
+        scope = scope.where(mount_key: mount_key) if mount_key
+        scope.first
+      end
     end
 
     def raw_code = @raw_code

data/app/models/toolchest/oauth_access_token.rb

@@ -18,6 +18,13 @@ module Toolchest
 
     def revoke! = update!(revoked_at: Time.current)
 
+    # Atomic revocation — returns true only if THIS call revoked the token.
+    # Prevents race conditions where two concurrent refresh exchanges both
+    # find the token active and both mint new token pairs.
+    def revoke_atomically!
+      self.class.where(id: id, revoked_at: nil).update_all(revoked_at: Time.current) > 0
+    end
+
     def accessible? = !revoked? && !expired?
 
     def scopes_array = (scopes || "").split(" ").reject(&:empty?)

data/lib/generators/toolchest/auth_generator.rb

@@ -19,7 +19,7 @@ module Toolchest
         end
       end
 
-      def create_migration
+      def copy_migration
         case strategy
         when "token"
           migration_template "create_toolchest_tokens.rb.tt",

data/lib/toolchest/app.rb

@@ -17,6 +17,12 @@ module Toolchest
     def call(env)
       Engine.ensure_initialized!
       env["toolchest.mount_key"] = @mount_key.to_s
+
+      cfg = Toolchest.configuration(@mount_key)
+      if cfg.mount_path.nil? && env["SCRIPT_NAME"].present?
+        cfg.mount_path = env["SCRIPT_NAME"]
+      end
+
       @router.call(env)
     end
 

data/lib/toolchest/auth/base.rb

@@ -7,7 +7,7 @@ module Toolchest
 
       def extract_bearer_token(request)
         auth_header = request.env["HTTP_AUTHORIZATION"] || ""
-        match = auth_header.match(/\ABearer\s+(.+)\z/i)
+        match = auth_header.match(/\ABearer\s+(\S+)\z/i)
         match&.[](1)
       end
     end

data/lib/toolchest/configuration.rb

@@ -7,7 +7,10 @@ module Toolchest
                   :scopes, :login_path, :additional_view_paths,
                   :access_token_expires_in, :toolboxes, :toolbox_module,
                   :mount_key, :mount_path,
-                  :optional_scopes, :required_scopes
+                  :optional_scopes, :required_scopes,
+                  :allowed_hosts, :allowed_origins, :dns_rebinding_protection,
+                  :session_idle_timeout, :max_sessions, :stateless,
+                  :page_size, :cache_ttl, :cache_scope
     attr_reader :auth
 
     def initialize(mount_key = :default)
@@ -28,6 +31,15 @@ module Toolchest
       @access_token_expires_in = 7200
       @toolboxes = nil
       @toolbox_module = nil
+      @allowed_hosts = nil
+      @allowed_origins = nil
+      @dns_rebinding_protection = nil
+      @session_idle_timeout = nil
+      @max_sessions = nil
+      @stateless = false
+      @page_size = nil
+      @cache_ttl = nil
+      @cache_scope = nil
     end
 
     def auth=(value)
@@ -92,5 +104,16 @@ module Toolchest
     end
 
     def resolved_server_name = @server_name || (defined?(Rails) && Rails.application ? Rails.application.class.module_parent_name : "Toolchest")
+
+    def transport_options
+      opts = {}
+      opts[:allowed_hosts] = @allowed_hosts if @allowed_hosts
+      opts[:allowed_origins] = @allowed_origins if @allowed_origins
+      opts[:dns_rebinding_protection] = @dns_rebinding_protection unless @dns_rebinding_protection.nil?
+      opts[:session_idle_timeout] = @session_idle_timeout if @session_idle_timeout
+      opts[:max_sessions] = @max_sessions if @max_sessions
+      opts[:stateless] = @stateless if @stateless
+      opts
+    end
   end
 end

data/lib/toolchest/current.rb

@@ -2,6 +2,6 @@ require "active_support/current_attributes"
 
 module Toolchest
   class Current < ActiveSupport::CurrentAttributes
-    attribute :auth, :mount_key, :mcp_session, :mcp_request_id, :mcp_progress_token
+    attribute :auth, :mount_key, :mcp_server_context, :mcp_session, :mcp_request_id, :mcp_progress_token
   end
 end

data/lib/toolchest/rack_app.rb

@@ -5,9 +5,8 @@ module Toolchest
     def initialize(mount_key: :default)
       @mount_key = mount_key.to_sym
       @server = build_mcp_server
-      @transport = MCP::Server::Transports::StreamableHTTPTransport.new(@server)
-      @server.transport = @transport
-      install_handlers!
+      # Transport auto-sets server.transport via super(server)
+      @transport = MCP::Server::Transports::StreamableHTTPTransport.new(@server, **config.transport_options)
     end
 
     def call(env)
@@ -49,9 +48,15 @@ module Toolchest
     def config = Toolchest.configuration(@mount_key)
 
     def build_mcp_server
+      router = Toolchest.router(@mount_key)
+
       opts = {
         name: config.resolved_server_name,
         version: config.server_version,
+        tools: build_mcp_tools(router),
+        prompts: build_mcp_prompts(router),
+        resources: build_mcp_resources(router),
+        resource_templates: build_mcp_resource_templates(router),
         capabilities: {
           tools: { listChanged: true },
           prompts: { listChanged: true },
@@ -63,43 +68,139 @@ module Toolchest
 
       opts[:description] = config.server_description if config.server_description
       opts[:instructions] = config.server_instructions if config.server_instructions
+      opts[:page_size] = config.page_size if config.page_size
+      opts[:ttl_ms] = config.cache_ttl if config.cache_ttl
+      opts[:cache_scope] = config.cache_scope if config.cache_scope
 
-      MCP::Server.new(**opts)
+      server = MCP::Server.new(**opts)
+      router.mcp_server = server
+      install_custom_handlers!(server, router)
+      server
     end
 
-    def install_handlers!
-      router = Toolchest.router(@mount_key)
-      server = @server
-
-      router.mcp_server = server
+    # Wire handlers that need custom dispatch logic:
+    # - tools/list: scope filtering based on current auth
+    # - resources/read: delegates to our router's resource lookup
+    # - completion/complete: delegates to our router's enum completion
+    def install_custom_handlers!(server, router)
+      # tools/list goes through the handler hash (else branch in dispatch),
+      # so overriding it gives us scope-filtered listing.
+      server.instance_variable_get(:@handlers)[MCP::Methods::TOOLS_LIST] = ->(params) {
+        { tools: router.tools_for_handler }
+      }
 
-      handlers = server.instance_variable_get(:@handlers)
-
-      handlers[MCP::Methods::TOOLS_LIST] = ->(params) { router.tools_for_handler }
-      handlers[MCP::Methods::RESOURCES_LIST] = ->(params) { router.resources_for_handler }
-      handlers[MCP::Methods::RESOURCES_READ] = ->(params) { router.resources_read_response(params) }
-      handlers[MCP::Methods::RESOURCES_TEMPLATES_LIST] = ->(params) { router.resource_templates_for_handler }
-      handlers[MCP::Methods::PROMPTS_LIST] = ->(params) { router.prompts_for_handler }
-      handlers[MCP::Methods::PROMPTS_GET] = ->(params) { router.prompts_get_response(params) }
-
-      # tools/call is hardcoded in handle_request to call private call_tool
-      server.define_singleton_method(:call_tool) do |params, session: nil, related_request_id: nil|
-        progress_token = params.dig(:_meta, :progressToken)
-        Toolchest::Current.mcp_session = session
-        Toolchest::Current.mcp_request_id = related_request_id
-        Toolchest::Current.mcp_progress_token = progress_token
-        router.dispatch_response(params)
+      # resources/read has a public setter — use it
+      server.resources_read_handler do |params|
+        router.resources_read_response(params)
       end
 
-      # completion/complete is hardcoded to call private complete, which validates
-      # against registered prompts/resources (we don't register any). override it.
-      server.define_singleton_method(:complete) do |params|
+      # completion/complete has a public setter — use it
+      server.completion_handler do |params|
         arg_name = params.dig(:argument, :name) || params.dig(:argument, "name")
         values = arg_name ? router.completion_values(arg_name) : []
         { completion: { values: values, hasMore: false } }
       end
     end
 
+    # --- Tool bridge ---
+    # Build an MCP::Tool subclass for each Toolchest::ToolDefinition.
+    # The tool's .call sets up Current context and delegates to Router#dispatch.
+
+    def build_mcp_tools(router)
+      naming = config.tool_naming
+      router.toolbox_classes.flat_map { |klass|
+        klass.tool_definitions.values.map { |td|
+          build_mcp_tool(td, router, naming)
+        }
+      }
+    end
+
+    def build_mcp_tool(td, router, naming)
+      tn = td.tool_name(naming)
+      desc = td.description
+      ttl = td.title
+      schema = td.input_schema
+      out_schema = td.output_schema
+      hints = td.resolved_annotations
+      r = router
+
+      Class.new(MCP::Tool) do
+        tool_name tn
+        title ttl if ttl
+        description desc
+        input_schema schema
+        output_schema out_schema if out_schema
+
+        if hints.any?
+          annotations({
+            read_only_hint: hints[:readOnlyHint],
+            destructive_hint: hints[:destructiveHint],
+            idempotent_hint: hints[:idempotentHint],
+            open_world_hint: hints[:openWorldHint],
+          }.compact)
+        end
+
+        define_singleton_method(:call) do |server_context: nil, **args|
+          Toolchest::Current.mcp_server_context = server_context
+          response = r.dispatch(tn, args)
+          MCP::Tool::Response.new(response[:content], error: response[:isError] || false)
+        end
+      end
+    end
+
+    # --- Prompt bridge ---
+    # Build an MCP::Prompt subclass for each prompt defined in toolboxes.
+    # The prompt's .template delegates to Router#prompts_get.
+
+    def build_mcp_prompts(router)
+      router.toolbox_classes.flat_map { |klass|
+        klass.prompts.map { |p| build_mcp_prompt(p, router) }
+      }
+    end
+
+    def build_mcp_prompt(pdef, router)
+      pn = pdef[:name]
+      pt = pdef[:title]
+      desc = pdef[:description]
+      args_list = (pdef[:arguments] || {}).map { |name, opts|
+        MCP::Prompt::Argument.new(
+          name: name.to_s,
+          description: opts[:description],
+          required: opts[:required] || false
+        )
+      }
+      r = router
+
+      Class.new(MCP::Prompt) do
+        prompt_name pn
+        title pt if pt
+        description desc
+        arguments args_list
+
+        define_singleton_method(:template) do |args, server_context: nil|
+          Toolchest::Current.mcp_server_context = server_context
+          r.prompts_get(pn, args || {})
+        end
+      end
+    end
+
+    # --- Resource bridge ---
+    # Create MCP::Resource instances for listing (reads go through resources_read_handler).
+
+    def build_mcp_resources(router)
+      router.toolbox_classes.flat_map(&:resources)
+        .reject { |r| r[:template] }
+        .map { |r| MCP::Resource.new(uri: r[:uri], name: r[:name], title: r[:title], description: r[:description]) }
+    end
+
+    def build_mcp_resource_templates(router)
+      router.toolbox_classes.flat_map(&:resources)
+        .select { |r| r[:template] }
+        .map { |r| MCP::ResourceTemplate.new(uri_template: r[:uri], name: r[:name], title: r[:title], description: r[:description]) }
+    end
+
+    # --- Auth ---
+
     def authenticate(request)
       strategy = case config.auth
       when :none then Auth::None.new

data/lib/toolchest/router.rb

@@ -60,6 +60,13 @@ module Toolchest
         end
       end
 
+      payload = {
+        tool: tool_name,
+        toolbox: definition.toolbox_class.name,
+        action: definition.method_name,
+        arguments: arguments
+      }
+
       start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       auth = Toolchest::Current.auth
       token_hint = extract_token_hint(auth)
@@ -75,6 +82,10 @@ module Toolchest
       duration = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time) * 1000).round(1)
       log_request_complete(definition, response, duration)
 
+      payload[:duration] = duration
+      payload[:error] = response[:isError] || false
+      ActiveSupport::Notifications.instrument("dispatch.toolchest", payload)
+
       response
     end
 

data/lib/toolchest/tool_definition.rb

@@ -1,8 +1,8 @@
 module Toolchest
   class ToolDefinition
-    attr_reader :method_name, :description, :params, :toolbox_class, :custom_name, :access_level, :scope, :annotations
+    attr_reader :method_name, :description, :params, :toolbox_class, :custom_name, :access_level, :scope, :annotations, :title, :output_schema
 
-    def initialize(method_name:, description:, params:, toolbox_class:, custom_name: nil, access_level: nil, scope: nil, annotations: nil)
+    def initialize(method_name:, description:, params:, toolbox_class:, custom_name: nil, access_level: nil, scope: nil, annotations: nil, title: nil, output_schema: nil)
       @method_name = method_name.to_sym
       @description = description
       @params = params
@@ -11,6 +11,8 @@ module Toolchest
       @access_level = access_level
       @scope = scope ? Array(scope) : nil
       @annotations = annotations
+      @title = title
+      @output_schema = output_schema
     end
 
     def tool_name(naming_strategy = nil)
@@ -25,6 +27,8 @@ module Toolchest
         description: @description,
         inputSchema: input_schema
       }
+      schema[:title] = @title if @title
+      schema[:outputSchema] = @output_schema if @output_schema
       hints = resolved_annotations
       schema[:annotations] = hints if hints.any?
       schema

data/lib/toolchest/toolbox.rb

@@ -48,10 +48,10 @@ module Toolchest
           .flat_map { |a| a.send(:own_prompts) }
       end
 
-      def tool(description, name: nil, access: nil, scope: nil, annotations: nil, &block)
+      def tool(description, name: nil, title: nil, access: nil, scope: nil, annotations: nil, output: nil, &block)
         builder = ToolBuilder.new
         builder.instance_eval(&block) if block
-        @_pending_tool = { description:, custom_name: name, access_level: access, scope:, annotations:, builder: }
+        @_pending_tool = { description:, custom_name: name, title:, access_level: access, scope:, annotations:, output_schema: output, builder: }
       end
 
       def default_param(name, type, description = "", **options)
@@ -67,11 +67,12 @@ module Toolchest
         }
       end
 
-      def resource(uri, name: nil, description: nil, &block)
+      def resource(uri, name: nil, title: nil, description: nil, &block)
         template = uri.include?("{")
         @_resources << {
           uri: uri,
           name: name || uri,
+          title: title,
           description: description,
           block: block,
           template: template,
@@ -79,9 +80,10 @@ module Toolchest
         }
       end
 
-      def prompt(prompt_name, description: nil, arguments: {}, &block)
+      def prompt(prompt_name, title: nil, description: nil, arguments: {}, &block)
         @_prompts << {
           name: prompt_name,
+          title: title,
           description: description,
           arguments: arguments,
           block: block,
@@ -111,9 +113,11 @@ module Toolchest
           params: params,
           toolbox_class: self,
           custom_name: pending[:custom_name],
+          title: pending[:title],
           access_level: pending[:access_level],
           scope: pending[:scope],
-          annotations: pending[:annotations]
+          annotations: pending[:annotations],
+          output_schema: pending[:output_schema]
         )
 
         @_tool_definitions[method_name.to_sym] = definition
@@ -229,17 +233,29 @@ module Toolchest
       throw :halt
     end
 
-    def mcp_log(level, message) = Toolchest.router(Toolchest::Current.mount_key&.to_sym || :default).notify_log(level: level.to_s, message: message)
+    def mcp_log(level, message)
+      ctx = Toolchest::Current.mcp_server_context
+      if ctx
+        ctx.notify_log_message(data: message, level: level.to_s, logger: "Toolchest")
+      else
+        Toolchest.router(Toolchest::Current.mount_key&.to_sym || :default).notify_log(level: level.to_s, message: message)
+      end
+    end
 
     # Report progress during long-running actions.
     # Client shows a progress bar. total and message are optional.
     def mcp_progress(progress, total: nil, message: nil)
+      ctx = Toolchest::Current.mcp_server_context
+      if ctx
+        ctx.report_progress(progress, total: total, message: message)
+        return
+      end
+
+      # Legacy fallback
       session = Toolchest::Current.mcp_session
       return unless session
-
       token = Toolchest::Current.mcp_progress_token
       return unless token
-
       session.notify_progress(
         progress_token: token,
         progress: progress,
@@ -260,7 +276,8 @@ module Toolchest
     #     s.temperature 0.3
     #   end
     def mcp_sample(prompt = nil, context: nil, max_tokens: 1024, **kwargs, &block)
-      session = Toolchest::Current.mcp_session
+      ctx = Toolchest::Current.mcp_server_context
+      session = ctx || Toolchest::Current.mcp_session
       raise Toolchest::Error, "Sampling requires an MCP client that supports it" unless session
 
       if block
@@ -282,11 +299,10 @@ module Toolchest
       end
 
       begin
-        result = session.create_sampling_message(
-          messages: messages,
-          related_request_id: Toolchest::Current.mcp_request_id,
-          **options
-        )
+        call_opts = { messages: messages, **options }
+        # Legacy path: explicitly pass related_request_id (ServerContext handles this automatically)
+        call_opts[:related_request_id] = Toolchest::Current.mcp_request_id unless ctx
+        result = session.create_sampling_message(**call_opts)
       rescue RuntimeError => e
         raise Toolchest::Error, "Sampling failed: #{e.message}"
       end
@@ -301,6 +317,43 @@ module Toolchest
       end
     end
 
+    # Ask the client's user for input via a form.
+    # Returns the elicitation result hash from the client.
+    #
+    #   result = mcp_elicit("Please confirm the shipping address",
+    #     schema: {
+    #       type: "object",
+    #       properties: {
+    #         confirmed: { type: "boolean", description: "Address is correct?" }
+    #       },
+    #       required: ["confirmed"]
+    #     })
+    #   halt error: "Cancelled" unless result["action"] == "accept"
+    #
+    def mcp_elicit(message, schema:)
+      ctx = Toolchest::Current.mcp_server_context
+      raise Toolchest::Error, "Elicitation requires an MCP client that supports it" unless ctx
+
+      begin
+        ctx.create_form_elicitation(message: message, requested_schema: schema)
+      rescue RuntimeError => e
+        raise Toolchest::Error, "Elicitation failed: #{e.message}"
+      end
+    end
+
+    # Check if the current request has been cancelled by the client.
+    def mcp_cancelled?
+      ctx = Toolchest::Current.mcp_server_context
+      ctx&.cancelled? || false
+    end
+
+    # Raise MCP::CancelledError if the current request has been cancelled.
+    # Use in long-running loops to bail early.
+    def mcp_raise_if_cancelled!
+      ctx = Toolchest::Current.mcp_server_context
+      ctx&.raise_if_cancelled!
+    end
+
     def dispatch(action_name)
       @_action_name = action_name
 

data/lib/toolchest/version.rb

@@ -1,3 +1,3 @@
 module Toolchest
-  VERSION = "0.3.6"
+  VERSION = "0.4.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: toolchest
 version: !ruby/object:Gem::Version
-  version: 0.3.6
+  version: 0.4.0
 platform: ruby
 authors:
 - Nora
@@ -29,14 +29,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0.10'
+        version: '1.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0.10'
+        version: '1.0'
 description: A Rails engine that maps the Model Context Protocol (MCP) to Rails conventions.
   If you've built a controller, you already know how to build a toolbox.
 executables: []