Add notification support
Co-authored-by: Kevin Fischer <kfischer_okarin@yahoo.co.jp>
This commit is contained in:
Riley Shott
parent
8fbb2e86aa
commit
55b03fc7aa
3 changed files with 711 additions and 6 deletions
|
|
@ -0,0 +1,249 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require "test_helper"
|
||||
require "json"
|
||||
require "mcp/server"
|
||||
require "mcp/server/transports/stdio_transport"
|
||||
|
||||
module MCP
|
||||
class Server
|
||||
module Transports
|
||||
class StdioNotificationIntegrationTest < ActiveSupport::TestCase
|
||||
class MockIO
|
||||
attr_reader :output
|
||||
|
||||
def initialize
|
||||
@output = []
|
||||
@closed = false
|
||||
end
|
||||
|
||||
def puts(message)
|
||||
@output << message
|
||||
end
|
||||
|
||||
def write(message)
|
||||
@output << message
|
||||
message.length
|
||||
end
|
||||
|
||||
def gets
|
||||
nil # Simulate end of input
|
||||
end
|
||||
|
||||
def set_encoding(encoding)
|
||||
# Mock implementation
|
||||
end
|
||||
|
||||
def flush
|
||||
# Mock implementation
|
||||
end
|
||||
|
||||
def close
|
||||
@closed = true
|
||||
end
|
||||
|
||||
def closed?
|
||||
@closed
|
||||
end
|
||||
end
|
||||
|
||||
setup do
|
||||
@original_stdout = $stdout
|
||||
@original_stdin = $stdin
|
||||
|
||||
@mock_stdout = MockIO.new
|
||||
@mock_stdin = MockIO.new
|
||||
|
||||
$stdout = @mock_stdout
|
||||
$stdin = @mock_stdin
|
||||
|
||||
@server = Server.new(
|
||||
name: "test_server",
|
||||
tools: [],
|
||||
prompts: [],
|
||||
resources: [],
|
||||
)
|
||||
@transport = StdioTransport.new(@server)
|
||||
@server.transport = @transport
|
||||
end
|
||||
|
||||
teardown do
|
||||
$stdout = @original_stdout
|
||||
$stdin = @original_stdin
|
||||
end
|
||||
|
||||
test "server notification methods send JSON-RPC notifications through StdioTransport" do
|
||||
# Test tools notification
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Test prompts notification
|
||||
@server.notify_prompts_list_changed
|
||||
|
||||
# Test resources notification
|
||||
@server.notify_resources_list_changed
|
||||
|
||||
# Check the notifications were sent
|
||||
assert_equal 3, @mock_stdout.output.size
|
||||
|
||||
# Parse and verify each notification
|
||||
notifications = @mock_stdout.output.map { |msg| JSON.parse(msg) }
|
||||
|
||||
assert_equal "2.0", notifications[0]["jsonrpc"]
|
||||
assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notifications[0]["method"]
|
||||
assert_nil notifications[0]["params"]
|
||||
|
||||
assert_equal "2.0", notifications[1]["jsonrpc"]
|
||||
assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notifications[1]["method"]
|
||||
assert_nil notifications[1]["params"]
|
||||
|
||||
assert_equal "2.0", notifications[2]["jsonrpc"]
|
||||
assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notifications[2]["method"]
|
||||
assert_nil notifications[2]["params"]
|
||||
end
|
||||
|
||||
test "notifications include params when provided" do
|
||||
# Test the transport's send_notification directly with params
|
||||
result = @transport.send_notification("test/notification", { data: "test_value" })
|
||||
|
||||
assert result
|
||||
assert_equal 1, @mock_stdout.output.size
|
||||
|
||||
notification = JSON.parse(@mock_stdout.output.first)
|
||||
assert_equal "2.0", notification["jsonrpc"]
|
||||
assert_equal "test/notification", notification["method"]
|
||||
assert_equal({ "data" => "test_value" }, notification["params"])
|
||||
end
|
||||
|
||||
test "server continues to work when stdout is closed" do
|
||||
# Close stdout
|
||||
@mock_stdout.close
|
||||
|
||||
# Server notifications should not raise errors
|
||||
assert_nothing_raised do
|
||||
@server.notify_tools_list_changed
|
||||
@server.notify_prompts_list_changed
|
||||
@server.notify_resources_list_changed
|
||||
end
|
||||
end
|
||||
|
||||
test "notifications work with dynamic tool additions" do
|
||||
# Define a new tool
|
||||
@server.define_tool(
|
||||
name: "dynamic_tool",
|
||||
description: "A dynamically added tool",
|
||||
) do |**args|
|
||||
{ result: "success" }
|
||||
end
|
||||
|
||||
# Clear previous output
|
||||
@mock_stdout.output.clear
|
||||
|
||||
# Manually trigger notification
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Check the notification was sent
|
||||
assert_equal 1, @mock_stdout.output.size
|
||||
|
||||
notification = JSON.parse(@mock_stdout.output.first)
|
||||
assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification["method"]
|
||||
|
||||
# Verify the tool was added to the server
|
||||
assert @server.tools.key?("dynamic_tool")
|
||||
end
|
||||
|
||||
test "notifications are properly formatted JSON-RPC 2.0 messages" do
|
||||
# Send a notification
|
||||
@server.notify_prompts_list_changed
|
||||
|
||||
# Verify format
|
||||
assert_equal 1, @mock_stdout.output.size
|
||||
output = @mock_stdout.output.first
|
||||
|
||||
# Should be valid JSON
|
||||
notification = JSON.parse(output)
|
||||
|
||||
# Should have required JSON-RPC 2.0 fields
|
||||
assert_equal "2.0", notification["jsonrpc"]
|
||||
assert notification.key?("method")
|
||||
refute notification.key?("id") # Notifications should not have an id
|
||||
|
||||
# Method should be the expected notification type
|
||||
assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification["method"]
|
||||
end
|
||||
|
||||
test "multiple notifications are sent as separate JSON messages" do
|
||||
# Send multiple notifications rapidly
|
||||
5.times do
|
||||
@server.notify_tools_list_changed
|
||||
end
|
||||
|
||||
# Each should be a separate JSON message
|
||||
assert_equal 5, @mock_stdout.output.size
|
||||
|
||||
# All should be parseable as JSON
|
||||
@mock_stdout.output.each do |msg|
|
||||
notification = JSON.parse(msg)
|
||||
assert_equal "2.0", notification["jsonrpc"]
|
||||
assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification["method"]
|
||||
end
|
||||
end
|
||||
|
||||
test "transport handles errors gracefully" do
|
||||
# Create a stdout that raises errors
|
||||
error_stdout = Class.new(MockIO) do
|
||||
def puts(message)
|
||||
raise IOError, "Simulated IO error"
|
||||
end
|
||||
end.new
|
||||
|
||||
$stdout = error_stdout
|
||||
|
||||
# Notification should return false but not raise
|
||||
result = @transport.send_notification("test/notification")
|
||||
refute result
|
||||
end
|
||||
|
||||
test "server notification flow works end-to-end with StdioTransport" do
|
||||
# This test verifies the complete integration from server to transport
|
||||
|
||||
# Start with no output
|
||||
assert_empty @mock_stdout.output
|
||||
|
||||
# Add a prompt and notify
|
||||
@server.define_prompt(
|
||||
name: "test_prompt",
|
||||
description: "Test prompt",
|
||||
) do |args, server_context:|
|
||||
MCP::PromptResponse.new(messages: [{ role: "user", content: "Test" }])
|
||||
end
|
||||
|
||||
# Manually trigger notification
|
||||
@server.notify_prompts_list_changed
|
||||
|
||||
# Verify notification was sent
|
||||
assert_equal 1, @mock_stdout.output.size
|
||||
notification = JSON.parse(@mock_stdout.output.first)
|
||||
assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification["method"]
|
||||
|
||||
# Add a resource and notify
|
||||
@server.resources = [
|
||||
MCP::Resource.new(
|
||||
uri: "test://resource",
|
||||
name: "Test Resource",
|
||||
description: "A test resource",
|
||||
mime_type: "text/plain",
|
||||
),
|
||||
]
|
||||
|
||||
# Manually trigger notification
|
||||
@server.notify_resources_list_changed
|
||||
|
||||
# Verify both notifications were sent
|
||||
assert_equal 2, @mock_stdout.output.size
|
||||
second_notification = JSON.parse(@mock_stdout.output.last)
|
||||
assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, second_notification["method"]
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
|
@ -0,0 +1,249 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require "test_helper"
|
||||
require "json"
|
||||
require "rack"
|
||||
require "mcp/server"
|
||||
require "mcp/server/transports/streamable_http_transport"
|
||||
|
||||
module MCP
|
||||
class Server
|
||||
module Transports
|
||||
class StreamableHTTPNotificationIntegrationTest < ActiveSupport::TestCase
|
||||
setup do
|
||||
@server = Server.new(
|
||||
name: "test_server",
|
||||
tools: [],
|
||||
prompts: [],
|
||||
resources: [],
|
||||
)
|
||||
@transport = StreamableHTTPTransport.new(@server)
|
||||
@server.transport = @transport
|
||||
end
|
||||
|
||||
test "server notification methods send SSE notifications through HTTP transport" do
|
||||
# Initialize a session
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
|
||||
)
|
||||
init_response = @transport.handle_request(init_request)
|
||||
session_id = init_response[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect with SSE
|
||||
io = StringIO.new
|
||||
get_request = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id },
|
||||
)
|
||||
response = @transport.handle_request(get_request)
|
||||
response[2].call(io) if response[2].is_a?(Proc)
|
||||
|
||||
# Give the stream time to set up
|
||||
sleep(0.1)
|
||||
|
||||
# Test tools notification
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Test prompts notification
|
||||
@server.notify_prompts_list_changed
|
||||
|
||||
# Test resources notification
|
||||
@server.notify_resources_list_changed
|
||||
|
||||
# Check the notifications were received
|
||||
io.rewind
|
||||
output = io.read
|
||||
|
||||
assert_includes output, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
|
||||
assert_includes output, "data: #{Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED}"
|
||||
assert_includes output, "data: #{Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED}"
|
||||
end
|
||||
|
||||
test "notifications are broadcast to all connected sessions" do
|
||||
# Create two sessions
|
||||
init_request1 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "123" }.to_json,
|
||||
)
|
||||
init_response1 = @transport.handle_request(init_request1)
|
||||
session_id1 = init_response1[1]["Mcp-Session-Id"]
|
||||
|
||||
init_request2 = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "456" }.to_json,
|
||||
)
|
||||
init_response2 = @transport.handle_request(init_request2)
|
||||
session_id2 = init_response2[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect both sessions with SSE
|
||||
io1 = StringIO.new
|
||||
get_request1 = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id1 },
|
||||
)
|
||||
response1 = @transport.handle_request(get_request1)
|
||||
response1[2].call(io1) if response1[2].is_a?(Proc)
|
||||
|
||||
io2 = StringIO.new
|
||||
get_request2 = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id2 },
|
||||
)
|
||||
response2 = @transport.handle_request(get_request2)
|
||||
response2[2].call(io2) if response2[2].is_a?(Proc)
|
||||
|
||||
# Give the streams time to set up
|
||||
sleep(0.1)
|
||||
|
||||
# Send notification through server
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Check both sessions received the notification
|
||||
io1.rewind
|
||||
output1 = io1.read
|
||||
assert_includes output1, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
|
||||
|
||||
io2.rewind
|
||||
output2 = io2.read
|
||||
assert_includes output2, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
|
||||
end
|
||||
|
||||
test "server continues to work when SSE connection is closed" do
|
||||
# Initialize a session
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
|
||||
)
|
||||
init_response = @transport.handle_request(init_request)
|
||||
session_id = init_response[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect with SSE
|
||||
io = StringIO.new
|
||||
get_request = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id },
|
||||
)
|
||||
response = @transport.handle_request(get_request)
|
||||
response[2].call(io) if response[2].is_a?(Proc)
|
||||
|
||||
# Give the stream time to set up
|
||||
sleep(0.1)
|
||||
|
||||
# Close the stream
|
||||
io.close
|
||||
|
||||
# Server notifications should not raise errors
|
||||
assert_nothing_raised do
|
||||
@server.notify_tools_list_changed
|
||||
@server.notify_prompts_list_changed
|
||||
@server.notify_resources_list_changed
|
||||
end
|
||||
end
|
||||
|
||||
test "notifications work with dynamic tool additions" do
|
||||
# Initialize a session
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
|
||||
)
|
||||
init_response = @transport.handle_request(init_request)
|
||||
session_id = init_response[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect with SSE
|
||||
io = StringIO.new
|
||||
get_request = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id },
|
||||
)
|
||||
response = @transport.handle_request(get_request)
|
||||
response[2].call(io) if response[2].is_a?(Proc)
|
||||
|
||||
# Give the stream time to set up
|
||||
sleep(0.1)
|
||||
|
||||
# Define a new tool (simulating dynamic tool addition)
|
||||
@server.define_tool(
|
||||
name: "dynamic_tool",
|
||||
description: "A dynamically added tool",
|
||||
) do |**args|
|
||||
{ result: "success" }
|
||||
end
|
||||
|
||||
# Manually trigger notification (since we removed the automatic triggers)
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Check the notification was received
|
||||
io.rewind
|
||||
output = io.read
|
||||
assert_includes output, "data: #{Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED}"
|
||||
|
||||
# Verify the tool was added to the server
|
||||
assert @server.tools.key?("dynamic_tool")
|
||||
end
|
||||
|
||||
test "SSE format is correct for notifications" do
|
||||
# Initialize a session
|
||||
init_request = create_rack_request(
|
||||
"POST",
|
||||
"/",
|
||||
{ "CONTENT_TYPE" => "application/json" },
|
||||
{ jsonrpc: "2.0", method: "initialize", id: "init" }.to_json,
|
||||
)
|
||||
init_response = @transport.handle_request(init_request)
|
||||
session_id = init_response[1]["Mcp-Session-Id"]
|
||||
|
||||
# Connect with SSE
|
||||
io = StringIO.new
|
||||
get_request = create_rack_request(
|
||||
"GET",
|
||||
"/",
|
||||
{ "HTTP_MCP_SESSION_ID" => session_id },
|
||||
)
|
||||
response = @transport.handle_request(get_request)
|
||||
response[2].call(io) if response[2].is_a?(Proc)
|
||||
|
||||
# Give the stream time to set up
|
||||
sleep(0.1)
|
||||
|
||||
# Send a notification
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
# Check SSE format
|
||||
io.rewind
|
||||
output = io.read
|
||||
|
||||
# SSE format should be "data: <message>\n\n"
|
||||
assert_match(/data: #{Regexp.escape(Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED)}\n/, output)
|
||||
end
|
||||
|
||||
private
|
||||
|
||||
def create_rack_request(method, path, headers, body = nil)
|
||||
env = {
|
||||
"REQUEST_METHOD" => method,
|
||||
"PATH_INFO" => path,
|
||||
"rack.input" => StringIO.new(body.to_s),
|
||||
}.merge(headers)
|
||||
|
||||
Rack::Request.new(env)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
136
test/mcp/server_notification_test.rb
Normal file
136
test/mcp/server_notification_test.rb
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
# typed: true
|
||||
# frozen_string_literal: true
|
||||
|
||||
require "test_helper"
|
||||
|
||||
module MCP
|
||||
class ServerNotificationTest < ActiveSupport::TestCase
|
||||
include InstrumentationTestHelper
|
||||
|
||||
class MockTransport < Transport
|
||||
attr_reader :notifications
|
||||
|
||||
def initialize(server)
|
||||
super
|
||||
@notifications = []
|
||||
end
|
||||
|
||||
def send_notification(method, params = nil)
|
||||
@notifications << { method: method, params: params }
|
||||
true
|
||||
end
|
||||
|
||||
def send_response(response); end
|
||||
def open; end
|
||||
def close; end
|
||||
def handle_request(request); end
|
||||
end
|
||||
|
||||
setup do
|
||||
configuration = MCP::Configuration.new
|
||||
configuration.instrumentation_callback = instrumentation_helper.callback
|
||||
|
||||
@server = Server.new(
|
||||
name: "test_server",
|
||||
version: "1.0.0",
|
||||
configuration: configuration
|
||||
)
|
||||
|
||||
@mock_transport = MockTransport.new(@server)
|
||||
@server.transport = @mock_transport
|
||||
end
|
||||
|
||||
test "#notify_tools_list_changed sends notification through transport" do
|
||||
@server.notify_tools_list_changed
|
||||
|
||||
assert_equal 1, @mock_transport.notifications.size
|
||||
notification = @mock_transport.notifications.first
|
||||
assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notification[:method]
|
||||
assert_nil notification[:params]
|
||||
end
|
||||
|
||||
test "#notify_prompts_list_changed sends notification through transport" do
|
||||
@server.notify_prompts_list_changed
|
||||
|
||||
assert_equal 1, @mock_transport.notifications.size
|
||||
notification = @mock_transport.notifications.first
|
||||
assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notification[:method]
|
||||
assert_nil notification[:params]
|
||||
end
|
||||
|
||||
test "#notify_resources_list_changed sends notification through transport" do
|
||||
@server.notify_resources_list_changed
|
||||
|
||||
assert_equal 1, @mock_transport.notifications.size
|
||||
notification = @mock_transport.notifications.first
|
||||
assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notification[:method]
|
||||
assert_nil notification[:params]
|
||||
end
|
||||
|
||||
test "notification methods work without transport" do
|
||||
server_without_transport = Server.new(name: "test_server")
|
||||
|
||||
# Should not raise any errors
|
||||
assert_nothing_raised do
|
||||
server_without_transport.notify_tools_list_changed
|
||||
server_without_transport.notify_prompts_list_changed
|
||||
server_without_transport.notify_resources_list_changed
|
||||
end
|
||||
end
|
||||
|
||||
test "notification methods handle transport errors gracefully" do
|
||||
# Create a transport that raises errors
|
||||
error_transport = Class.new(MockTransport) do
|
||||
def send_notification(method, params = nil)
|
||||
raise StandardError, "Transport error"
|
||||
end
|
||||
end.new(@server)
|
||||
|
||||
@server.transport = error_transport
|
||||
|
||||
# Mock the exception reporter
|
||||
expected_contexts = [
|
||||
{ notification: "tools_list_changed" },
|
||||
{ notification: "prompts_list_changed" },
|
||||
{ notification: "resources_list_changed" }
|
||||
]
|
||||
|
||||
call_count = 0
|
||||
@server.configuration.exception_reporter.expects(:call).times(3).with do |exception, context|
|
||||
assert_kind_of StandardError, exception
|
||||
assert_equal "Transport error", exception.message
|
||||
assert_includes expected_contexts, context
|
||||
call_count += 1
|
||||
true
|
||||
end
|
||||
|
||||
# Should not raise errors to the caller
|
||||
assert_nothing_raised do
|
||||
@server.notify_tools_list_changed
|
||||
@server.notify_prompts_list_changed
|
||||
@server.notify_resources_list_changed
|
||||
end
|
||||
|
||||
assert_equal 3, call_count
|
||||
end
|
||||
|
||||
test "multiple notification methods can be called in sequence" do
|
||||
@server.notify_tools_list_changed
|
||||
@server.notify_prompts_list_changed
|
||||
@server.notify_resources_list_changed
|
||||
|
||||
assert_equal 3, @mock_transport.notifications.size
|
||||
|
||||
notifications = @mock_transport.notifications
|
||||
assert_equal Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED, notifications[0][:method]
|
||||
assert_equal Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED, notifications[1][:method]
|
||||
assert_equal Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED, notifications[2][:method]
|
||||
end
|
||||
|
||||
test "notification method constants are defined correctly" do
|
||||
assert_equal "notifications/tools/list_changed", Methods::NOTIFICATIONS_TOOLS_LIST_CHANGED
|
||||
assert_equal "notifications/prompts/list_changed", Methods::NOTIFICATIONS_PROMPTS_LIST_CHANGED
|
||||
assert_equal "notifications/resources/list_changed", Methods::NOTIFICATIONS_RESOURCES_LIST_CHANGED
|
||||
end
|
||||
end
|
||||
end
|
||||
Loading…
Reference in a new issue