streamable_http_server.rb

# frozen_string_literal: true

# Usage: bundle exec ruby examples/streamable_http_server.rb
require "mcp"
require "rack/cors"
require "rackup"
require "json"
require "logger"

# Create a logger for SSE-specific logging
sse_logger = Logger.new($stdout)
sse_logger.formatter = proc do |severity, datetime, _progname, msg|
  "[SSE] #{severity} #{datetime.strftime("%H:%M:%S.%L")} - #{msg}\n"
end

# Tool that returns a response that will be sent via SSE if a stream is active
class NotificationTool < MCP::Tool
  tool_name "notification_tool"
  description "Returns a notification message that will be sent via SSE if stream is active"
  input_schema(
    properties: {
      message: { type: "string", description: "Message to send via SSE" },
      delay: { type: "number", description: "Delay in seconds before returning (optional)" },
    },
    required: ["message"],
  )

  class << self
    attr_accessor :logger

    def call(message:, delay: 0)
      sleep(delay) if delay > 0

      logger.info("Returning notification message: #{message}")

      MCP::Tool::Response.new([{
        type: "text",
        text: "Notification: #{message} (timestamp: #{Time.now.iso8601})",
      }])
    end
  end
end

# Create the server
server = MCP::Server.new(
  name: "sse_test_server",
  tools: [NotificationTool],
  prompts: [],
  resources: [],
)

# Set logger for tools
NotificationTool.logger = sse_logger

# Add a simple echo tool for basic testing
server.define_tool(
  name: "echo",
  description: "Simple echo tool",
  input_schema: { properties: { message: { type: "string" } }, required: ["message"] },
) do |message:|
  MCP::Tool::Response.new([{ type: "text", text: "Echo: #{message}" }])
end

# Create the Streamable HTTP transport
transport = MCP::Server::Transports::StreamableHTTPTransport.new(server)
server.transport = transport

# Create a logger for MCP request/response logging
mcp_logger = Logger.new($stdout)
mcp_logger.formatter = proc do |_severity, _datetime, _progname, msg|
  "[MCP] #{msg}\n"
end

# Create the Rack application
app = proc do |env|
  request = Rack::Request.new(env)

  # Log request details
  if request.post?
    body = request.body.read
    request.body.rewind
    begin
      parsed_body = JSON.parse(body)
      mcp_logger.info("Request: #{parsed_body["method"]} (id: #{parsed_body["id"]})")

      # Log SSE-specific setup
      if parsed_body["method"] == "initialize"
        sse_logger.info("New client initializing session")
      end
    rescue JSON::ParserError
      mcp_logger.warn("Invalid JSON in request")
    end
  elsif request.get?
    session_id = request.env["HTTP_MCP_SESSION_ID"] ||
      Rack::Utils.parse_query(request.env["QUERY_STRING"])["sessionId"]
    sse_logger.info("SSE connection request for session: #{session_id}")
  end

  # Handle the request
  response = transport.handle_request(request)

  # Log response details
  status, headers, body = response
  if body.is_a?(Array) && !body.empty? && request.post?
    begin
      parsed_response = JSON.parse(body.first)
      if parsed_response["error"]
        mcp_logger.error("Response error: #{parsed_response["error"]["message"]}")
      elsif parsed_response["accepted"]
        # Response was sent via SSE
        server.notify_log_message(data: { details: "Response accepted and sent via SSE" }, level: "info")
        sse_logger.info("Response sent via SSE stream")
      else
        mcp_logger.info("Response: success (id: #{parsed_response["id"]})")

        # Log session ID for initialization
        if headers["Mcp-Session-Id"]
          sse_logger.info("Session created: #{headers["Mcp-Session-Id"]}")
        end
      end
    rescue JSON::ParserError
      mcp_logger.warn("Invalid JSON in response")
    end
  elsif request.get? && status == 200
    sse_logger.info("SSE stream established")
  end

  response
end

# Build the Rack application with middleware
rack_app = Rack::Builder.new do
  # Enable CORS to allow browser-based MCP clients (e.g., MCP Inspector)
  # WARNING: origins("*") allows all origins. Restrict this in production.
  use(Rack::Cors) do
    allow do
      origins("*")
      resource(
        "*",
        headers: :any,
        methods: [:get, :post, :delete, :options],
        expose: ["Mcp-Session-Id"],
      )
    end
  end

  use(Rack::CommonLogger, Logger.new($stdout))
  use(Rack::ShowExceptions)
  run(app)
end

# Print usage instructions
puts <<~MESSAGE
  === MCP Streaming HTTP Test Server ===

  Starting server on http://localhost:9393

  Available Tools:
  1. NotificationTool - Returns messages that are sent via SSE when stream is active"
  2. echo - Simple echo tool

  Testing SSE:

  1. Initialize session:
     curl -i http://localhost:9393 \\
       --json '{"jsonrpc":"2.0","method":"initialize","id":1,"params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"sse-test","version":"1.0"}}}'

  2. Connect SSE stream (use the session ID from step 1):"
     curl -i -N -H "Mcp-Session-Id: YOUR_SESSION_ID" http://localhost:9393

  3. In another terminal, test tools (responses will be sent via SSE if stream is active):

     Echo tool:
     curl -i http://localhost:9393 -H "Mcp-Session-Id: YOUR_SESSION_ID" \\
       --json '{"jsonrpc":"2.0","method":"tools/call","id":2,"params":{"name":"echo","arguments":{"message":"Hello SSE!"}}}'

     Notification tool (with 2 second delay):
     curl -i http://localhost:9393 -H "Mcp-Session-Id: YOUR_SESSION_ID" \\
       --json '{"jsonrpc":"2.0","method":"tools/call","id":3,"params":{"name":"notification_tool","arguments":{"message":"Hello SSE!", "delay": 2}}}'

  Note: When an SSE stream is active, tool responses will appear in the SSE stream and the POST request will return {"accepted": true}

  Press Ctrl+C to stop the server
MESSAGE

# Start the server
Rackup::Handler.get("puma").run(rack_app, Port: 9393, Host: "localhost")