ruby-sdk/test/mcp/server_context_test.rb
Koichi ITO da5f9225ac Support sampling/createMessage per MCP specification
## Motivation and Context

The MCP specification defines `sampling/createMessage` to request LLM completions from the client.
The Ruby SDK did not yet support server-to-client requests,
so the server had no way to ask the client to sample an LLM.
This caused the `tools-call-sampling` conformance test to fail.

This adds `Server#create_sampling_message` to send a `sampling/createMessage` request
to the client via the transport layer. The method validates that the client declared
the `sampling` capability during initialization, and optionally checks the
`sampling.tools` capability when tools or tool_choice parameters are provided.

On the transport side, `Transport#send_request` is introduced as an abstract method,
with implementations in both `StdioTransport` and `StreamableHTTPTransport`.
Each transport assigns a unique request ID, sends the JSON-RPC request to the client,
and blocks until the client responds.

Client capabilities are stored per session via `ServerSession`, so validation is scoped
to the originating client. `ServerSession#create_sampling_message` uses
`build_sampling_params` for capability validation and sends the request directly through
the transport with the session ID, following the same pattern as `notify_progress`.

`Server#create_sampling_message` is the public API for single-client transports
(e.g., StdioTransport).

For multi-client transports (e.g., StreamableHTTPTransport),
`ServerSession#create_sampling_message` must be used to route requests to the correct client.
Tools using `server_context.create_sampling_message` automatically route to the correct session
regardless of transport.

`StreamableHTTPTransport#send_request` requires `session_id` to prevent broadcasting sampling requests
to all connected clients. Missing or invalid session IDs raise explicit errors instead of silently failing.

The `include_context` parameter is soft-deprecated in the MCP specification but is included
for compatibility with existing clients that declare the `sampling.context` capability.
Python and TypeScript SDKs also retain this parameter.

Ref: https://modelcontextprotocol.io/specification/latest/server/utilities/sampling

## How Has This Been Tested?

Added comprehensive tests in `test/mcp/server_sampling_test.rb` covering required and optional parameters,
capability validation, error handling, nil-param omission, per-session capability isolation,
and HTTP init capability scoping.

Added transport tests in `test/mcp/server/transports/stdio_transport_test.rb` and
`test/mcp/server/transports/streamable_http_transport_test.rb` for `send_request` round-trip behavior,
error responses, stateless mode rejection, missing session_id, invalid session_id, and no-active-stream error.

Conformance: `tools-call-sampling` is removed from expected failures. The conformance server's `TestSampling`
tool now calls `server_context.create_sampling_message` instead of returning a stub error.

This aligns the Ruby SDK with the MCP specification and other SDK implementations.

## Breaking Changes

None for end users.
`Transport#send_request` is a new abstract method. `ServerSession#create_sampling_message` is a new method.
2026-04-04 05:55:27 +09:00

784 lines
24 KiB
Ruby

# frozen_string_literal: true
require "test_helper"
module MCP
class ServerContextTest < ActiveSupport::TestCase
test "ServerContext delegates method calls to the underlying context" do
context = { user: "test_user" }
progress = Progress.new(notification_target: mock, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
assert_equal "test_user", server_context[:user]
end
test "ServerContext respond_to? returns true for methods on the underlying context" do
context = { user: "test_user" }
progress = Progress.new(notification_target: mock, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
assert_respond_to server_context, :[]
assert_respond_to server_context, :keys
end
test "ServerContext respond_to? returns false for methods not on the underlying context" do
context = { user: "test_user" }
progress = Progress.new(notification_target: mock, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
refute_respond_to server_context, :nonexistent_method
end
test "ServerContext raises NoMethodError for methods not on the underlying context" do
context = { user: "test_user" }
progress = Progress.new(notification_target: mock, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
assert_raises(NoMethodError) { server_context.nonexistent_method }
end
test "ServerContext#create_sampling_message delegates to notification_target over context" do
notification_target = mock
notification_target.expects(:create_sampling_message).with(
messages: [{ role: "user", content: { type: "text", text: "Hello" } }],
max_tokens: 100,
).returns({ role: "assistant", content: { type: "text", text: "Hi" } })
context = mock
progress = Progress.new(notification_target: notification_target, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: notification_target)
result = server_context.create_sampling_message(
messages: [{ role: "user", content: { type: "text", text: "Hello" } }],
max_tokens: 100,
)
assert_equal "Hi", result[:content][:text]
end
test "ServerContext#create_sampling_message falls back to context when notification_target does not respond" do
notification_target = mock
context = mock
context.expects(:create_sampling_message).with(
messages: [{ role: "user", content: { type: "text", text: "Hello" } }],
max_tokens: 100,
).returns({ role: "assistant", content: { type: "text", text: "Fallback" } })
progress = Progress.new(notification_target: notification_target, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: notification_target)
result = server_context.create_sampling_message(
messages: [{ role: "user", content: { type: "text", text: "Hello" } }],
max_tokens: 100,
)
assert_equal "Fallback", result[:content][:text]
end
test "ServerContext delegates to custom object context" do
context = Object.new
def context.custom_method
"custom_value"
end
progress = Progress.new(notification_target: mock, progress_token: nil)
server_context = ServerContext.new(context, progress: progress, notification_target: mock)
assert_equal "custom_value", server_context.custom_method
end
test "ServerContext#report_progress works with nil context" do
progress = mock
progress.expects(:report).with(50, total: 100, message: nil).once
server_context = ServerContext.new(nil, progress: progress, notification_target: mock)
server_context.report_progress(50, total: 100)
end
test "ServerContext#notify_log_message is a no-op when notification_target is nil" do
progress = Progress.new(notification_target: nil, progress_token: nil)
server_context = ServerContext.new(nil, progress: progress, notification_target: nil)
assert_nothing_raised { server_context.notify_log_message(data: "test", level: "info") }
end
# Tool without server_context parameter
class SimpleToolWithoutContext < Tool
tool_name "simple_without_context"
description "A tool that doesn't use server_context"
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
class << self
def call(message:)
Tool::Response.new([
{ type: "text", content: "SimpleToolWithoutContext: #{message}" },
])
end
end
end
# Tool with optional server_context parameter
class ToolWithOptionalContext < Tool
tool_name "tool_with_optional_context"
description "A tool with optional server_context"
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
class << self
def call(message:, server_context: nil)
user = server_context.respond_to?(:[]) ? server_context[:user] : nil
context_info = user ? "with context: #{user}" : "no context"
Tool::Response.new([
{ type: "text", content: "ToolWithOptionalContext: #{message} (#{context_info})" },
])
end
end
end
# Tool with required server_context parameter
class ToolWithRequiredContext < Tool
tool_name "tool_with_required_context"
description "A tool that requires server_context"
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
class << self
def call(message:, server_context:)
Tool::Response.new([
{ type: "text", content: "ToolWithRequiredContext: #{message} for user #{server_context[:user]}" },
])
end
end
end
setup do
@server_with_context = Server.new(
name: "test_server",
tools: [SimpleToolWithoutContext, ToolWithOptionalContext, ToolWithRequiredContext],
server_context: { user: "test_user" },
)
@server_without_context = Server.new(
name: "test_server_no_context",
tools: [SimpleToolWithoutContext, ToolWithOptionalContext],
)
end
test "tool without server_context parameter works when server has context" do
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "simple_without_context",
arguments: { message: "Hello" },
},
}
response = @server_with_context.handle(request)
assert response[:result]
assert_equal "SimpleToolWithoutContext: Hello", response[:result][:content][0][:content]
end
test "tool with optional server_context receives context when server has it" do
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_optional_context",
arguments: { message: "Hello" },
},
}
response = @server_with_context.handle(request)
assert response[:result]
assert_equal "ToolWithOptionalContext: Hello (with context: test_user)",
response[:result][:content][0][:content]
end
test "tool with optional server_context works when server has no context" do
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_optional_context",
arguments: { message: "Hello" },
},
}
response = @server_without_context.handle(request)
assert response[:result]
assert_equal "ToolWithOptionalContext: Hello (no context)",
response[:result][:content][0][:content]
end
test "tool with required server_context receives context" do
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_required_context",
arguments: { message: "Hello" },
},
}
response = @server_with_context.handle(request)
assert response[:result]
assert_equal "ToolWithRequiredContext: Hello for user test_user",
response[:result][:content][0][:content]
end
test "tool with required server_context returns protocol error in JSON-RPC format when server has no context" do
server_no_context = Server.new(
name: "test_server_no_context",
tools: [ToolWithRequiredContext],
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_required_context",
arguments: { message: "Hello" },
},
}
response = server_no_context.handle(request)
assert_nil response[:result]
assert_equal(-32603, response[:error][:code])
assert_equal "Internal error", response[:error][:message]
assert_match(/Internal error calling tool tool_with_required_context: /, response[:error][:data])
end
test "call_tool_with_args correctly detects server_context parameter presence" do
# Tool without server_context
refute SimpleToolWithoutContext.method(:call).parameters.any? { |_type, name| name == :server_context }
# Tool with optional server_context
assert ToolWithOptionalContext.method(:call).parameters.any? { |_type, name| name == :server_context }
# Tool with required server_context
assert ToolWithRequiredContext.method(:call).parameters.any? { |_type, name| name == :server_context }
end
test "tools can use splat kwargs to accept any arguments including server_context" do
class FlexibleTool < Tool
tool_name "flexible_tool"
class << self
def call(**kwargs)
message = kwargs[:message]
context = kwargs[:server_context]
Tool::Response.new([
{
type: "text",
content: "FlexibleTool: #{message} (context: #{context ? "present" : "absent"})",
},
])
end
end
end
server = Server.new(
name: "test_server",
tools: [FlexibleTool],
server_context: { user: "test_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "flexible_tool",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "FlexibleTool: Hello (context: present)",
response[:result][:content][0][:content]
end
# Prompt tests
# Prompt without server_context parameter
class SimplePromptWithoutContext < Prompt
prompt_name "simple_prompt_without_context"
description "A prompt that doesn't use server_context"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args)
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new("SimplePromptWithoutContext: #{args[:message]}"),
),
],
)
end
end
end
# Prompt with optional server_context parameter
class PromptWithOptionalContext < Prompt
prompt_name "prompt_with_optional_context"
description "A prompt with optional server_context"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args, server_context: nil)
user = server_context.respond_to?(:[]) ? server_context[:user] : nil
context_info = user ? "with context: #{user}" : "no context"
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new("PromptWithOptionalContext: #{args[:message]} (#{context_info})"),
),
],
)
end
end
end
# Prompt with required server_context parameter
class PromptWithRequiredContext < Prompt
prompt_name "prompt_with_required_context"
description "A prompt that requires server_context"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args, server_context:)
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new(
"PromptWithRequiredContext: #{args[:message]} for user #{server_context[:user]}",
),
),
],
)
end
end
end
test "prompt without server_context parameter works when server has context" do
server = Server.new(
name: "test_server",
prompts: [SimplePromptWithoutContext],
server_context: { user: "test_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "simple_prompt_without_context",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "SimplePromptWithoutContext: Hello", response[:result][:messages][0][:content][:text]
end
test "prompt with optional server_context receives context when server has it" do
server = Server.new(
name: "test_server",
prompts: [PromptWithOptionalContext],
server_context: { user: "test_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "prompt_with_optional_context",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "PromptWithOptionalContext: Hello (with context: test_user)",
response[:result][:messages][0][:content][:text]
end
test "prompt with optional server_context works when server has no context" do
server = Server.new(
name: "test_server",
prompts: [PromptWithOptionalContext],
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "prompt_with_optional_context",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "PromptWithOptionalContext: Hello (no context)",
response[:result][:messages][0][:content][:text]
end
test "prompt with required server_context receives context" do
server = Server.new(
name: "test_server",
prompts: [PromptWithRequiredContext],
server_context: { user: "test_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "prompt_with_required_context",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "PromptWithRequiredContext: Hello for user test_user",
response[:result][:messages][0][:content][:text]
end
test "prompts can use splat kwargs to accept any arguments including server_context" do
class FlexiblePrompt < Prompt
prompt_name "flexible_prompt"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args, **kwargs)
message = args[:message]
context = kwargs[:server_context]
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new("FlexiblePrompt: #{message} (context: #{context ? "present" : "absent"})"),
),
],
)
end
end
end
server = Server.new(
name: "test_server",
prompts: [FlexiblePrompt],
server_context: { user: "test_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "flexible_prompt",
arguments: { message: "Hello" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "FlexiblePrompt: Hello (context: present)",
response[:result][:messages][0][:content][:text]
end
test "tool receives _meta when provided in request params" do
class ToolWithMeta < Tool
tool_name "tool_with_meta"
description "A tool that uses _meta"
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
class << self
def call(message:, server_context:)
meta_info = server_context.dig(:_meta, :provider, :metadata)
Tool::Response.new([
{ type: "text", content: "Message: #{message}, Metadata: #{meta_info}" },
])
end
end
end
server = Server.new(
name: "test_server",
tools: [ToolWithMeta],
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_meta",
arguments: { message: "Hello" },
_meta: {
provider: {
metadata: "test_value",
},
},
},
}
response = server.handle(request)
assert response[:result]
assert_equal "Message: Hello, Metadata: test_value",
response[:result][:content][0][:content]
end
test "_meta is nested within server_context" do
class ToolWithNestedMeta < Tool
tool_name "tool_with_nested_meta"
description "A tool that uses nested _meta"
input_schema({ properties: { message: { type: "string" } }, required: ["message"] })
class << self
def call(message:, server_context:)
user = server_context[:user]
session_id = server_context.dig(:_meta, :session_id)
Tool::Response.new([
{ type: "text", content: "User: #{user}, Session: #{session_id}, Message: #{message}" },
])
end
end
end
server = Server.new(
name: "test_server",
tools: [ToolWithNestedMeta],
server_context: { user: "test_user", original_field: "value" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_nested_meta",
arguments: { message: "Hello" },
_meta: {
session_id: "abc123",
},
},
}
response = server.handle(request)
assert response[:result]
assert_equal "User: test_user, Session: abc123, Message: Hello",
response[:result][:content][0][:content]
end
test "_meta preserves original server_context" do
class ToolPreservesContext < Tool
tool_name "tool_preserves_context"
description "A tool that checks context preservation"
class << self
def call(server_context:)
priority = server_context[:priority]
meta_priority = server_context.dig(:_meta, :priority)
Tool::Response.new([
{ type: "text", content: "Context priority: #{priority}, Meta priority: #{meta_priority}" },
])
end
end
end
server = Server.new(
name: "test_server",
tools: [ToolPreservesContext],
server_context: { priority: "low" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_preserves_context",
arguments: {},
_meta: {
priority: "high",
},
},
}
response = server.handle(request)
assert response[:result]
assert_equal "Context priority: low, Meta priority: high", response[:result][:content][0][:content]
end
test "prompt receives _meta when provided in request params" do
class PromptWithMeta < Prompt
prompt_name "prompt_with_meta"
description "A prompt that uses _meta"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args, server_context:)
meta_info = server_context.dig(:_meta, :request_id)
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new("Message: #{args[:message]}, Request ID: #{meta_info}"),
),
],
)
end
end
end
server = Server.new(
name: "test_server",
prompts: [PromptWithMeta],
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "prompt_with_meta",
arguments: { message: "Hello" },
_meta: {
request_id: "req_12345",
},
},
}
response = server.handle(request)
assert response[:result]
assert_equal "Message: Hello, Request ID: req_12345",
response[:result][:messages][0][:content][:text]
end
test "non-Hash server_context is preserved when _meta is present" do
server_obj = Object.new
def server_obj.custom_method
"custom_value"
end
tool_class = Class.new(Tool) do
tool_name "tool_with_non_hash_context"
define_singleton_method(:call) do |server_context:, **_args|
Tool::Response.new([
{ type: "text", content: "custom: #{server_context.custom_method}" },
])
end
end
server = Server.new(
name: "test_server",
tools: [tool_class],
server_context: server_obj,
)
request = {
jsonrpc: "2.0",
id: 1,
method: "tools/call",
params: {
name: "tool_with_non_hash_context",
arguments: {},
_meta: { progressToken: "token-1" },
},
}
response = server.handle(request)
assert response[:result]
assert_equal "custom: custom_value", response[:result][:content][0][:content]
end
test "_meta is nested within server_context for prompts" do
class PromptWithNestedContext < Prompt
prompt_name "prompt_with_nested_context"
description "A prompt that uses nested context"
arguments [Prompt::Argument.new(name: "message", required: true)]
class << self
def template(args, server_context:)
user = server_context[:user]
trace_id = server_context.dig(:_meta, :trace_id)
Prompt::Result.new(
messages: [
Prompt::Message.new(
role: "user",
content: Content::Text.new("User: #{user}, Trace: #{trace_id}, Message: #{args[:message]}"),
),
],
)
end
end
end
server = Server.new(
name: "test_server",
prompts: [PromptWithNestedContext],
server_context: { user: "prompt_user" },
)
request = {
jsonrpc: "2.0",
id: 1,
method: "prompts/get",
params: {
name: "prompt_with_nested_context",
arguments: { message: "World" },
_meta: {
trace_id: "trace_xyz789",
},
},
}
response = server.handle(request)
assert response[:result]
assert_equal "User: prompt_user, Trace: trace_xyz789, Message: World",
response[:result][:messages][0][:content][:text]
end
end
end