Merge pull request #332 from koic/feature_cancellation
Support `notifications/cancelled` per MCP specification
This commit is contained in:
commit
3d9437736b
9 changed files with 1485 additions and 41 deletions
|
|
@ -18,6 +18,11 @@ module JsonRpcHandler
|
|||
|
||||
DEFAULT_ALLOWED_ID_CHARACTERS = /\A[a-zA-Z0-9_-]+\z/
|
||||
|
||||
# Sentinel return value from a handler. When a handler returns this,
|
||||
# `process_request` emits no JSON-RPC response for the request,
|
||||
# matching the notification-style semantics (id is ignored).
|
||||
NO_RESPONSE = Object.new.freeze
|
||||
|
||||
extend self
|
||||
|
||||
def handle(request, id_validation_pattern: DEFAULT_ALLOWED_ID_CHARACTERS, &method_finder)
|
||||
|
|
@ -103,6 +108,7 @@ module JsonRpcHandler
|
|||
end
|
||||
|
||||
result = method.call(params)
|
||||
return if result.equal?(NO_RESPONSE)
|
||||
|
||||
success_response(id: id, result: result)
|
||||
rescue MCP::Server::RequestHandlerError => e
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ require_relative "mcp/version"
|
|||
|
||||
module MCP
|
||||
autoload :Annotations, "mcp/annotations"
|
||||
autoload :Cancellation, "mcp/cancellation"
|
||||
autoload :CancelledError, "mcp/cancelled_error"
|
||||
autoload :Client, "mcp/client"
|
||||
autoload :Content, "mcp/content"
|
||||
autoload :Icon, "mcp/icon"
|
||||
|
|
|
|||
72
lib/mcp/cancellation.rb
Normal file
72
lib/mcp/cancellation.rb
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require_relative "cancelled_error"
|
||||
|
||||
module MCP
|
||||
class Cancellation
|
||||
attr_reader :reason, :request_id
|
||||
|
||||
def initialize(request_id: nil)
|
||||
@request_id = request_id
|
||||
@reason = nil
|
||||
@cancelled = false
|
||||
@callbacks = []
|
||||
@mutex = Mutex.new
|
||||
end
|
||||
|
||||
def cancelled?
|
||||
@mutex.synchronize { @cancelled }
|
||||
end
|
||||
|
||||
def cancel(reason: nil)
|
||||
callbacks = @mutex.synchronize do
|
||||
return false if @cancelled
|
||||
|
||||
@cancelled = true
|
||||
@reason = reason
|
||||
@callbacks.tap { @callbacks = [] }
|
||||
end
|
||||
|
||||
callbacks.each do |callback|
|
||||
callback.call(reason)
|
||||
rescue StandardError => e
|
||||
MCP.configuration.exception_reporter.call(e, { error: "Cancellation callback failed" })
|
||||
end
|
||||
|
||||
true
|
||||
end
|
||||
|
||||
# Registers a callback invoked synchronously on the first `cancel` call.
|
||||
# If already cancelled, fires immediately.
|
||||
#
|
||||
# Returns the block itself as a handle that can be passed to `off_cancel`
|
||||
# to deregister it (e.g. when a nested request completes normally and the
|
||||
# hook should not fire on a later parent cancellation).
|
||||
def on_cancel(&block)
|
||||
fire_now = false
|
||||
@mutex.synchronize do
|
||||
if @cancelled
|
||||
fire_now = true
|
||||
else
|
||||
@callbacks << block
|
||||
end
|
||||
end
|
||||
|
||||
block.call(@reason) if fire_now
|
||||
block
|
||||
end
|
||||
|
||||
# Removes a previously-registered `on_cancel` callback. Returns `true`
|
||||
# if the callback was still pending (i.e. had not yet fired), `false`
|
||||
# otherwise. Safe to call with `nil`.
|
||||
def off_cancel(handle)
|
||||
return false unless handle
|
||||
|
||||
@mutex.synchronize { !@callbacks.delete(handle).nil? }
|
||||
end
|
||||
|
||||
def raise_if_cancelled!
|
||||
raise CancelledError.new(request_id: @request_id, reason: @reason) if cancelled?
|
||||
end
|
||||
end
|
||||
end
|
||||
13
lib/mcp/cancelled_error.rb
Normal file
13
lib/mcp/cancelled_error.rb
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
module MCP
|
||||
class CancelledError < StandardError
|
||||
attr_reader :request_id, :reason
|
||||
|
||||
def initialize(message = "Request was cancelled", request_id: nil, reason: nil)
|
||||
super(message)
|
||||
@request_id = request_id
|
||||
@reason = reason
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require_relative "../json_rpc_handler"
|
||||
require_relative "cancellation"
|
||||
require_relative "cancelled_error"
|
||||
require_relative "instrumentation"
|
||||
require_relative "methods"
|
||||
require_relative "logging_message_notification"
|
||||
|
|
@ -384,6 +386,13 @@ module MCP
|
|||
end
|
||||
|
||||
def handle_request(request, method, session: nil, related_request_id: nil)
|
||||
# `notifications/cancelled` is dispatched directly: it is a notification (no JSON-RPC id)
|
||||
# and intentionally bypasses the `@handlers` lookup, capability check, in-flight registry,
|
||||
# and rescue blocks below.
|
||||
if method == Methods::NOTIFICATIONS_CANCELLED
|
||||
return ->(params) { handle_cancelled_notification(params, session: session) }
|
||||
end
|
||||
|
||||
handler = @handlers[method]
|
||||
unless handler
|
||||
instrument_call("unsupported_method", server_context: { request: request }) do
|
||||
|
|
@ -395,6 +404,12 @@ module MCP
|
|||
|
||||
Methods.ensure_capability!(method, capabilities)
|
||||
|
||||
# `initialize` MUST NOT be cancelled (MCP spec 2025-11-25, cancellation item 2),
|
||||
# so do not track it in the in-flight registry.
|
||||
cancellation = if related_request_id && method != Methods::INITIALIZE
|
||||
session&.register_in_flight(related_request_id)
|
||||
end
|
||||
|
||||
->(params) {
|
||||
reported_exception = nil
|
||||
instrument_call(
|
||||
|
|
@ -406,23 +421,33 @@ module MCP
|
|||
when Methods::INITIALIZE
|
||||
init(params, session: session)
|
||||
when Methods::RESOURCES_READ
|
||||
{ contents: @handlers[Methods::RESOURCES_READ].call(params) }
|
||||
{ contents: read_resource_contents(params, session: session, related_request_id: related_request_id, cancellation: cancellation) }
|
||||
when Methods::RESOURCES_SUBSCRIBE, Methods::RESOURCES_UNSUBSCRIBE
|
||||
@handlers[method].call(params)
|
||||
dispatch_optional_context_handler(@handlers[method], params, session: session, related_request_id: related_request_id, cancellation: cancellation)
|
||||
{}
|
||||
when Methods::TOOLS_CALL
|
||||
call_tool(params, session: session, related_request_id: related_request_id)
|
||||
call_tool(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
|
||||
when Methods::PROMPTS_GET
|
||||
get_prompt(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
|
||||
when Methods::COMPLETION_COMPLETE
|
||||
complete(params)
|
||||
complete(params, session: session, related_request_id: related_request_id, cancellation: cancellation)
|
||||
when Methods::LOGGING_SET_LEVEL
|
||||
configure_logging_level(params, session: session)
|
||||
else
|
||||
@handlers[method].call(params)
|
||||
dispatch_optional_context_handler(@handlers[method], params, session: session, related_request_id: related_request_id, cancellation: cancellation)
|
||||
end
|
||||
client = session&.client || @client
|
||||
add_instrumentation_data(client: client) if client
|
||||
|
||||
if cancellation&.cancelled?
|
||||
add_instrumentation_data(cancelled: true, cancellation_reason: cancellation.reason)
|
||||
next JsonRpcHandler::NO_RESPONSE
|
||||
end
|
||||
|
||||
result
|
||||
rescue CancelledError => e
|
||||
add_instrumentation_data(cancelled: true, cancellation_reason: e.reason)
|
||||
next JsonRpcHandler::NO_RESPONSE
|
||||
rescue RequestHandlerError => e
|
||||
report_exception(e.original_error || e, { request: request })
|
||||
add_instrumentation_data(error: e.error_type)
|
||||
|
|
@ -434,10 +459,23 @@ module MCP
|
|||
wrapped = RequestHandlerError.new("Internal error handling #{method} request", request, original_error: e)
|
||||
reported_exception = wrapped
|
||||
raise wrapped
|
||||
ensure
|
||||
session&.unregister_in_flight(related_request_id) if related_request_id
|
||||
end
|
||||
}
|
||||
end
|
||||
|
||||
def handle_cancelled_notification(params, session: nil)
|
||||
return unless session
|
||||
return unless params.is_a?(Hash)
|
||||
|
||||
request_id = params[:requestId] || params["requestId"]
|
||||
return if request_id.nil?
|
||||
|
||||
reason = params[:reason] || params["reason"]
|
||||
session.cancel_incoming(request_id: request_id, reason: reason)
|
||||
end
|
||||
|
||||
def default_capabilities
|
||||
{
|
||||
tools: { listChanged: true },
|
||||
|
|
@ -516,7 +554,7 @@ module MCP
|
|||
{ tools: page[:items], nextCursor: page[:next_cursor] }.compact
|
||||
end
|
||||
|
||||
def call_tool(request, session: nil, related_request_id: nil)
|
||||
def call_tool(request, session: nil, related_request_id: nil, cancellation: nil)
|
||||
tool_name = request[:name]
|
||||
|
||||
tool = tools[tool_name]
|
||||
|
|
@ -548,8 +586,12 @@ module MCP
|
|||
|
||||
progress_token = request.dig(:_meta, :progressToken)
|
||||
|
||||
call_tool_with_args(tool, arguments, server_context_with_meta(request), progress_token: progress_token, session: session, related_request_id: related_request_id)
|
||||
rescue RequestHandlerError
|
||||
call_tool_with_args(
|
||||
tool, arguments, server_context_with_meta(request), progress_token: progress_token, session: session, related_request_id: related_request_id, cancellation: cancellation
|
||||
)
|
||||
rescue RequestHandlerError, CancelledError
|
||||
# CancelledError is intentionally not wrapped so `handle_request` can turn it into
|
||||
# `JsonRpcHandler::NO_RESPONSE` per the MCP cancellation spec.
|
||||
raise
|
||||
rescue => e
|
||||
raise RequestHandlerError.new(
|
||||
|
|
@ -566,7 +608,7 @@ module MCP
|
|||
{ prompts: page[:items], nextCursor: page[:next_cursor] }.compact
|
||||
end
|
||||
|
||||
def get_prompt(request)
|
||||
def get_prompt(request, session: nil, related_request_id: nil, cancellation: nil)
|
||||
prompt_name = request[:name]
|
||||
prompt = @prompts[prompt_name]
|
||||
unless prompt
|
||||
|
|
@ -579,7 +621,14 @@ module MCP
|
|||
prompt_args = request[:arguments]
|
||||
prompt.validate_arguments!(prompt_args)
|
||||
|
||||
call_prompt_template_with_args(prompt, prompt_args, server_context_with_meta(request))
|
||||
server_context = build_server_context(
|
||||
request: request,
|
||||
session: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
|
||||
call_prompt_template_with_args(prompt, prompt_args, server_context)
|
||||
end
|
||||
|
||||
def list_resources(request)
|
||||
|
|
@ -600,14 +649,82 @@ module MCP
|
|||
{ resourceTemplates: page[:items], nextCursor: page[:next_cursor] }.compact
|
||||
end
|
||||
|
||||
def complete(params)
|
||||
def complete(params, session: nil, related_request_id: nil, cancellation: nil)
|
||||
validate_completion_params!(params)
|
||||
|
||||
result = @handlers[Methods::COMPLETION_COMPLETE].call(params)
|
||||
result = dispatch_optional_context_handler(
|
||||
@handlers[Methods::COMPLETION_COMPLETE],
|
||||
params,
|
||||
session: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
|
||||
normalize_completion_result(result)
|
||||
end
|
||||
|
||||
# Invokes `resources/read` via the registered handler. If the handler block opts in to `server_context:`,
|
||||
# pass an `MCP::ServerContext` so the handler can observe cancellation via `server_context.cancelled?` or
|
||||
# `server_context.raise_if_cancelled!`.
|
||||
def read_resource_contents(request, session: nil, related_request_id: nil, cancellation: nil)
|
||||
dispatch_optional_context_handler(
|
||||
@handlers[Methods::RESOURCES_READ],
|
||||
request,
|
||||
session: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
end
|
||||
|
||||
# Opt-in `server_context:` dispatch for block-based handlers registered via `resources_read_handler`,
|
||||
# `completion_handler`, `resources_subscribe_handler`, `resources_unsubscribe_handler`, or `define_custom_method`.
|
||||
# Existing handlers that only accept `params` are called unchanged; handlers that declare a `server_context:`
|
||||
# keyword receive an `MCP::ServerContext` wrapping the raw server context with cancellation plumbing.
|
||||
def dispatch_optional_context_handler(handler, params, session: nil, related_request_id: nil, cancellation: nil)
|
||||
return handler.call(params) unless handler_declares_server_context?(handler)
|
||||
|
||||
server_context = build_server_context(
|
||||
request: params,
|
||||
session: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
handler.call(params, server_context: server_context)
|
||||
end
|
||||
|
||||
# Stricter than `accepts_server_context?`: requires `server_context` to appear as a named keyword parameter
|
||||
# (`:key` optional, `:keyreq` required). Positional parameters named `server_context` (`:req` / `:opt`) are NOT
|
||||
# treated as opt-in - otherwise `handler.call(params, server_context: ctx)` would pass the `{server_context: ctx}`
|
||||
# Hash as the handler's second positional argument, which is never what the user meant.
|
||||
#
|
||||
# `**kwargs`-only signatures (`:keyrest` without a named `server_context`) are also not opt-in here,
|
||||
# because the dispatch site passes a positional `params`, and a `**kwargs`-only block cannot accept
|
||||
# that positional argument (lambdas/methods raise `ArgumentError`; non-lambda procs silently drop `params`).
|
||||
# Tool handlers intentionally allow `**kwargs` opt-in via `accepts_server_context?` because they are invoked
|
||||
# via `tool.call(**args, server_context: …)` without a positional argument.
|
||||
def handler_declares_server_context?(handler)
|
||||
return false unless handler.respond_to?(:parameters)
|
||||
|
||||
handler.parameters.any? do |type, name|
|
||||
name == :server_context && (type == :key || type == :keyreq)
|
||||
end
|
||||
end
|
||||
|
||||
# Builds an `MCP::ServerContext` used to give a handler access to session-scoped helpers
|
||||
# (progress, cancellation, nested server-to-client requests).
|
||||
def build_server_context(request:, session:, related_request_id:, cancellation:)
|
||||
meta_source = request.is_a?(Hash) ? request : {}
|
||||
progress_token = meta_source.dig(:_meta, :progressToken)
|
||||
progress = Progress.new(notification_target: session, progress_token: progress_token, related_request_id: related_request_id)
|
||||
ServerContext.new(
|
||||
server_context_with_meta(meta_source),
|
||||
progress: progress,
|
||||
notification_target: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
end
|
||||
|
||||
def report_exception(exception, server_context = {})
|
||||
configuration.exception_reporter.call(exception, server_context)
|
||||
end
|
||||
|
|
@ -628,18 +745,32 @@ module MCP
|
|||
).to_h
|
||||
end
|
||||
|
||||
# Whether a tool/prompt handler opts in to receiving an `MCP::ServerContext`.
|
||||
# Recognizes `:keyrest` (`**kwargs`) because tools are invoked without a positional argument
|
||||
# (`tool.call(**args, server_context:)`), soa `**kwargs`-only signature safely captures `server_context:`.
|
||||
# Named keyword `server_context` must be `:key` or `:keyreq` - positional parameters (`:req` / `:opt`) that
|
||||
# happen to be named `server_context` are excluded because the call site passes `server_context:` as a keyword,
|
||||
# and a positional slot would receive the `{server_context: ctx}` Hash instead.
|
||||
def accepts_server_context?(method_object)
|
||||
parameters = method_object.parameters
|
||||
|
||||
parameters.any? { |type, name| type == :keyrest || name == :server_context }
|
||||
parameters.any? do |type, name|
|
||||
type == :keyrest || (name == :server_context && (type == :key || type == :keyreq))
|
||||
end
|
||||
end
|
||||
|
||||
def call_tool_with_args(tool, arguments, context, progress_token: nil, session: nil, related_request_id: nil)
|
||||
def call_tool_with_args(tool, arguments, context, progress_token: nil, session: nil, related_request_id: nil, cancellation: nil)
|
||||
args = arguments&.transform_keys(&:to_sym) || {}
|
||||
|
||||
if accepts_server_context?(tool.method(:call))
|
||||
progress = Progress.new(notification_target: session, progress_token: progress_token, related_request_id: related_request_id)
|
||||
server_context = ServerContext.new(context, progress: progress, notification_target: session, related_request_id: related_request_id)
|
||||
server_context = ServerContext.new(
|
||||
context,
|
||||
progress: progress,
|
||||
notification_target: session,
|
||||
related_request_id: related_request_id,
|
||||
cancellation: cancellation,
|
||||
)
|
||||
tool.call(**args, server_context: server_context).to_h
|
||||
else
|
||||
tool.call(**args).to_h
|
||||
|
|
|
|||
|
|
@ -54,6 +54,13 @@ module MCP
|
|||
false
|
||||
end
|
||||
|
||||
# NOTE: This signature deliberately matches the abstract `Transport#send_request` contract
|
||||
# (`method, params = nil`) without the cancellation kwargs that `StreamableHTTPTransport#send_request` accepts.
|
||||
# On Ruby 2.7 the project's supported minimum a method that mixes a positional `params` Hash with
|
||||
# explicit keyword arguments cannot be called as `send_request(method, { ... })` - the trailing Hash would be
|
||||
# auto-promoted to keyword arguments. Stdio is single-threaded and blocks on `$stdin.gets`, so nested-request
|
||||
# cancellation has very limited value here regardless; servers that need cancellation propagation for nested
|
||||
# server-to-client requests should use `StreamableHTTPTransport`.
|
||||
def send_request(method, params = nil)
|
||||
request_id = generate_request_id
|
||||
request = { jsonrpc: "2.0", id: request_id, method: method }
|
||||
|
|
|
|||
|
|
@ -175,7 +175,7 @@ module MCP
|
|||
# sends the request via SSE stream, then blocks on `queue.pop`.
|
||||
# When the client POSTs a response, `handle_response` matches it by `request_id`
|
||||
# and pushes the result onto the queue, unblocking this thread.
|
||||
def send_request(method, params = nil, session_id: nil, related_request_id: nil)
|
||||
def send_request(method, params = nil, session_id: nil, related_request_id: nil, parent_cancellation: nil, server_session: nil)
|
||||
if @stateless
|
||||
raise "Stateless mode does not support server-to-client requests."
|
||||
end
|
||||
|
|
@ -190,6 +190,7 @@ module MCP
|
|||
|
||||
request_id = generate_request_id
|
||||
queue = Queue.new
|
||||
cancel_hook = nil
|
||||
|
||||
request = { jsonrpc: "2.0", id: request_id, method: method }
|
||||
request[:params] = params if params
|
||||
|
|
@ -229,6 +230,16 @@ module MCP
|
|||
raise "No active stream for #{method} request."
|
||||
end
|
||||
|
||||
if parent_cancellation && server_session
|
||||
cancel_hook = parent_cancellation.on_cancel do |reason|
|
||||
server_session.send_peer_cancellation(
|
||||
nested_request_id: request_id,
|
||||
related_request_id: related_request_id,
|
||||
reason: reason,
|
||||
)
|
||||
end
|
||||
end
|
||||
|
||||
response = queue.pop
|
||||
|
||||
if response.is_a?(Hash) && response.key?(:error)
|
||||
|
|
@ -239,8 +250,18 @@ module MCP
|
|||
raise "SSE session closed while waiting for #{method} response."
|
||||
end
|
||||
|
||||
if response == :cancelled
|
||||
reason = @mutex.synchronize { @pending_responses.dig(request_id, :cancel_reason) }
|
||||
raise MCP::CancelledError.new(
|
||||
"#{method} request was cancelled",
|
||||
request_id: request_id,
|
||||
reason: reason,
|
||||
)
|
||||
end
|
||||
|
||||
response
|
||||
ensure
|
||||
parent_cancellation.off_cancel(cancel_hook) if cancel_hook
|
||||
if request_id
|
||||
@mutex.synchronize do
|
||||
@pending_responses.delete(request_id)
|
||||
|
|
@ -248,6 +269,24 @@ module MCP
|
|||
end
|
||||
end
|
||||
|
||||
# Unblocks a `send_request` awaiting a response when the peer is being cancelled.
|
||||
# The waiting thread will see `:cancelled` on its queue and raise `MCP::CancelledError`.
|
||||
#
|
||||
# Race note: this is first-writer-wins on the pending-response queue. If a real response
|
||||
# has already been pushed (client responded before the cancel hook fired), that response
|
||||
# wins and `:cancelled` is enqueued behind it but never read - `send_request` returns
|
||||
# the real response and deletes the pending entry in its `ensure` block. Conversely,
|
||||
# if `:cancelled` arrives first, any later client response is silently dropped in `handle_response`
|
||||
# because the pending entry has been removed.
|
||||
def cancel_pending_request(request_id, reason: nil)
|
||||
@mutex.synchronize do
|
||||
if (pending = @pending_responses[request_id])
|
||||
pending[:cancel_reason] = reason
|
||||
pending[:queue].push(:cancelled)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
private
|
||||
|
||||
def start_reaper_thread
|
||||
|
|
@ -309,6 +348,7 @@ module MCP
|
|||
return missing_session_id_response if !@stateless && !session_id
|
||||
|
||||
if notification?(body)
|
||||
dispatch_notification(body_string, session_id)
|
||||
handle_accepted
|
||||
elsif response?(body)
|
||||
return session_not_found_response if !@stateless && !session_exists?(session_id)
|
||||
|
|
@ -459,6 +499,22 @@ module MCP
|
|||
!body[:id] && !!body[:method]
|
||||
end
|
||||
|
||||
# Dispatches a client-originated notification (e.g. `notifications/cancelled`,
|
||||
# `notifications/initialized`) through the server so it can update session state.
|
||||
def dispatch_notification(body_string, session_id)
|
||||
server_session = nil
|
||||
if session_id && !@stateless
|
||||
@mutex.synchronize do
|
||||
session = @sessions[session_id]
|
||||
server_session = session[:server_session] if session
|
||||
end
|
||||
end
|
||||
|
||||
dispatch_handle_json(body_string, server_session)
|
||||
rescue => e
|
||||
MCP.configuration.exception_reporter.call(e, { error: "Failed to dispatch notification" })
|
||||
end
|
||||
|
||||
def response?(body)
|
||||
!!body[:id] && !body[:method]
|
||||
end
|
||||
|
|
@ -536,6 +592,12 @@ module MCP
|
|||
handle_request_with_sse_response(body_string, session_id, server_session, related_request_id: related_request_id)
|
||||
else
|
||||
response = dispatch_handle_json(body_string, server_session)
|
||||
|
||||
# `Server#handle_json` returns `nil` when cancellation has suppressed the JSON-RPC response per spec.
|
||||
# Mirror the notification path and ack with 202 instead of returning a 200 with a `nil` Rack body,
|
||||
# which would produce an empty body the client cannot parse as JSON.
|
||||
return handle_accepted if response.nil?
|
||||
|
||||
[200, { "Content-Type" => "application/json" }, [response]]
|
||||
end
|
||||
end
|
||||
|
|
|
|||
|
|
@ -2,11 +2,22 @@
|
|||
|
||||
module MCP
|
||||
class ServerContext
|
||||
def initialize(context, progress:, notification_target:, related_request_id: nil)
|
||||
attr_reader :cancellation
|
||||
|
||||
def initialize(context, progress:, notification_target:, related_request_id: nil, cancellation: nil)
|
||||
@context = context
|
||||
@progress = progress
|
||||
@notification_target = notification_target
|
||||
@related_request_id = related_request_id
|
||||
@cancellation = cancellation
|
||||
end
|
||||
|
||||
def cancelled?
|
||||
!!@cancellation&.cancelled?
|
||||
end
|
||||
|
||||
def raise_if_cancelled!
|
||||
@cancellation&.raise_if_cancelled!
|
||||
end
|
||||
|
||||
# Reports progress for the current tool operation.
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require_relative "cancellation"
|
||||
require_relative "methods"
|
||||
|
||||
module MCP
|
||||
|
|
@ -15,6 +16,48 @@ module MCP
|
|||
@client = nil
|
||||
@client_capabilities = nil
|
||||
@logging_message_notification = nil
|
||||
@in_flight = {}
|
||||
@in_flight_mutex = Mutex.new
|
||||
end
|
||||
|
||||
# Registers a `Cancellation` token for an in-flight request.
|
||||
def register_in_flight(request_id)
|
||||
return if request_id.nil?
|
||||
|
||||
cancellation = Cancellation.new(request_id: request_id)
|
||||
@in_flight_mutex.synchronize { @in_flight[request_id] = cancellation }
|
||||
cancellation
|
||||
end
|
||||
|
||||
def unregister_in_flight(request_id)
|
||||
return if request_id.nil?
|
||||
|
||||
@in_flight_mutex.synchronize { @in_flight.delete(request_id) }
|
||||
end
|
||||
|
||||
def lookup_in_flight(request_id)
|
||||
@in_flight_mutex.synchronize { @in_flight[request_id] }
|
||||
end
|
||||
|
||||
# Flips the `Cancellation` for a matching in-flight request received from the peer.
|
||||
# Silently ignores unknown IDs per MCP spec (cancellation utilities, item 5).
|
||||
def cancel_incoming(request_id:, reason: nil)
|
||||
cancellation = lookup_in_flight(request_id)
|
||||
cancellation&.cancel(reason: reason)
|
||||
end
|
||||
|
||||
# Sends `notifications/cancelled` to the peer for a previously-issued request.
|
||||
# Also unblocks any transport-level `send_request` waiting on a response for `request_id`.
|
||||
def cancel_request(request_id:, reason: nil)
|
||||
params = { requestId: request_id }
|
||||
params[:reason] = reason if reason
|
||||
send_to_transport(Methods::NOTIFICATIONS_CANCELLED, params)
|
||||
|
||||
if @transport.respond_to?(:cancel_pending_request)
|
||||
@transport.cancel_pending_request(request_id, reason: reason)
|
||||
end
|
||||
rescue => e
|
||||
MCP.configuration.exception_reporter.call(e, { notification: "cancelled", request_id: request_id })
|
||||
end
|
||||
|
||||
def handle(request)
|
||||
|
|
@ -78,6 +121,23 @@ module MCP
|
|||
send_to_transport_request(Methods::ELICITATION_CREATE, params, related_request_id: related_request_id)
|
||||
end
|
||||
|
||||
# Sends `notifications/cancelled` to the peer for a nested server-to-client request
|
||||
# that was started inside a now-cancelled parent request. `related_request_id`
|
||||
# is the parent request id so the notification is routed to the same stream
|
||||
# (e.g. the parent's POST response stream on `StreamableHTTPTransport`) rather than
|
||||
# the GET SSE stream.
|
||||
def send_peer_cancellation(nested_request_id:, related_request_id: nil, reason: nil)
|
||||
params = { requestId: nested_request_id }
|
||||
params[:reason] = reason if reason
|
||||
send_to_transport(Methods::NOTIFICATIONS_CANCELLED, params, related_request_id: related_request_id)
|
||||
|
||||
if @transport.respond_to?(:cancel_pending_request)
|
||||
@transport.cancel_pending_request(nested_request_id, reason: reason)
|
||||
end
|
||||
rescue => e
|
||||
MCP.configuration.exception_reporter.call(e, { notification: "cancelled", request_id: nested_request_id })
|
||||
end
|
||||
|
||||
# Sends an elicitation complete notification scoped to this session.
|
||||
def notify_elicitation_complete(elicitation_id:)
|
||||
send_to_transport(Methods::NOTIFICATIONS_ELICITATION_COMPLETE, { elicitationId: elicitation_id })
|
||||
|
|
@ -121,32 +181,57 @@ module MCP
|
|||
|
||||
private
|
||||
|
||||
# Branches on `@session_id` because `StdioTransport` creates a `ServerSession` without
|
||||
# a `session_id` (`session_id: nil`), while `StreamableHTTPTransport` always provides one.
|
||||
#
|
||||
# TODO: When Ruby 2.7 support is dropped, replace with a direct call:
|
||||
# `@transport.send_notification(method, params, session_id: @session_id)` and
|
||||
# add `**` to `Transport#send_notification` and `StdioTransport#send_notification`.
|
||||
# Forwards `send_notification` to the transport with only the kwargs the transport's method signature
|
||||
# actually accepts. Custom transports that implement the abstract `send_notification(method, params = nil)`
|
||||
# contract continue to work unchanged; bundled transports that declare `session_id:` / `related_request_id:`
|
||||
# receive the session-scoped routing information.
|
||||
def send_to_transport(method, params, related_request_id: nil)
|
||||
if @session_id
|
||||
@transport.send_notification(method, params, session_id: @session_id, related_request_id: related_request_id)
|
||||
else
|
||||
@transport.send_notification(method, params)
|
||||
end
|
||||
kwargs = {
|
||||
session_id: @session_id,
|
||||
related_request_id: related_request_id,
|
||||
}.compact
|
||||
|
||||
forward_to_transport(@transport.method(:send_notification), method, params, kwargs)
|
||||
end
|
||||
|
||||
# Branches on `@session_id` because `StdioTransport` creates a `ServerSession` without
|
||||
# a `session_id` (`session_id: nil`), while `StreamableHTTPTransport` always provides one.
|
||||
#
|
||||
# TODO: When Ruby 2.7 support is dropped, replace with a direct call:
|
||||
# `@transport.send_request(method, params, session_id: @session_id)` and
|
||||
# add `**` to `Transport#send_request` and `StdioTransport#send_request`.
|
||||
# Forwards `send_request` to the transport with only the kwargs the transport's method signature
|
||||
# actually accepts. Custom transports that implement the abstract `send_request(method, params = nil)`
|
||||
# contract continue to work; bundled transports that declare `session_id:` / `related_request_id:` /
|
||||
# `parent_cancellation:` / `server_session:` receive the nested-cancellation plumbing.
|
||||
# When `related_request_id` names an in-flight request, its `Cancellation` token is looked up
|
||||
# so that cancelling the parent also cancels this nested server-to-client request.
|
||||
def send_to_transport_request(method, params, related_request_id: nil)
|
||||
if @session_id
|
||||
@transport.send_request(method, params, session_id: @session_id, related_request_id: related_request_id)
|
||||
parent_cancellation = related_request_id ? lookup_in_flight(related_request_id) : nil
|
||||
|
||||
kwargs = {
|
||||
session_id: @session_id,
|
||||
related_request_id: related_request_id,
|
||||
parent_cancellation: parent_cancellation,
|
||||
server_session: self,
|
||||
}.compact
|
||||
|
||||
forward_to_transport(@transport.method(:send_request), method, params, kwargs)
|
||||
end
|
||||
|
||||
# Calls `transport_method(method, params, **supported)` where `supported` contains only the keys
|
||||
# the transport's method signature accepts. This keeps bundled transports (which declare the new kwargs)
|
||||
# working while preserving compatibility with custom transports that implement only the abstract
|
||||
# `(method, params = nil)` contract.
|
||||
def forward_to_transport(transport_method, method, params, kwargs)
|
||||
parameters = transport_method.parameters
|
||||
accepts_keyrest = parameters.any? { |type, _| type == :keyrest }
|
||||
supported = if accepts_keyrest
|
||||
kwargs
|
||||
else
|
||||
@transport.send_request(method, params)
|
||||
allowed = parameters.filter_map { |type, name| name if type == :key || type == :keyreq }
|
||||
kwargs.slice(*allowed)
|
||||
end
|
||||
|
||||
# Always splat `**supported` even when empty: on Ruby 2.7 the bare `transport_method.call(method, params)`
|
||||
# form would let the trailing `params` Hash be auto-promoted to keyword arguments when the receiver
|
||||
# accepts `**kwargs`, breaking handlers that rely on `params` arriving as a positional Hash.
|
||||
# The explicit splat suppresses that conversion and is a no-op when `supported` is empty.
|
||||
transport_method.call(method, params, **supported)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
|
|||
Loading…
Reference in a new issue