model-context-protocol-rb 0.5.1 → 0.6.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8fd0681c65b47196375ec7f4934cadf226ed01b1b289caf296348698808ef5cb
-  data.tar.gz: 4ce234855d56b724b6f69b97851ff1ee4db65370b7361ea0a65c43a929c448b7
+  metadata.gz: f0f2318890bd52a6ade4d17652273ed887c62cc2badab13aabeb79a993c34a11
+  data.tar.gz: d85af9411a3ebb83ca4938c61cd586d9511152b11aea9b91e0765493fd39508b
 SHA512:
-  metadata.gz: b3bafa3fe92aa05ff380ad5fc311c0143d5e1bb6c3a91e50893115eb01ad2ea8b2d93dde89240815b50629fcb3b3b3467585766ac42d2c7087cd0be55b431455
-  data.tar.gz: 9d1c4f1c5f93352bf9392f9aeaf63253534348787aecb0866abe160dbf7cebc302bfec751955b19d0a06d9d60675c450960d401af2538aa9b100fe076bc7fd94
+  metadata.gz: 6557e13bb133a661f26de2d53e6b8b82ca044f26ed2ec288cc03b297699b0cef4de7d540014bbf345317c1d456cc3b27986df146b52604922636e80f8eae5bcb
+  data.tar.gz: 065604a28cc5b21bf690d7c138a2be3015ba1d862db7e7393dbcead3ef94f86d54c70e0d7abad51d21f7b60aa65c6b49a54143851326c57a71813a186b698c9c

data/CHANGELOG.md

@@ -1,5 +1,16 @@
 ## [Unreleased]
 
+## [0.6.0] - 2025-01-26
+
+- Implement server logging capability for internal server diagnostics.
+- Add informational and debug logging in streamable HTTP transport.
+- (Fix) Ensure stream monitor thread shuts down quickly.
+- (Fix) Ensure streams are closed when the server shuts down.
+- (Fix) Fix stream handling in streamable HTTP transport.
+- (Fix) Ensure empty options don't break the registry.
+- (Fix) Ensure progressable timer tasks do not run indefinitely.
+- (Breaking) Update connection pool dependency; requires updating other gems that depend on connection_pool.
+
 ## [0.5.1] - 2025-09-23
 
 - (Fix) Ensure streams are properly closed when clients disconnect.
@@ -78,7 +89,8 @@
 
 - Initial release
 
-[Unreleased]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.5.0...HEAD
+[Unreleased]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.6.0...HEAD
+[0.6.0]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.5.1...v0.6.0
 [0.5.1]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.5.0...v0.5.1
 [0.5.0]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.4.0...v0.5.0
 [0.4.0]: https://github.com/dickdavis/model-context-protocol-rb/compare/v0.3.4...v0.4.0

data/README.md

@@ -14,6 +14,7 @@ Provides simple abstractions that allow you to serve prompts, resources, resourc
   - [Pagination Configuration Options](#pagination-configuration-options)
   - [Transport Configuration Options](#transport-configuration-options)
   - [Redis Configuration](#redis-configuration)
+  - [Server Logging Configuration](#server-logging-configuration)
   - [Registry Configuration Options](#registry-configuration-options)
 - [Prompts](#prompts)
 - [Resources](#resources)
@@ -91,7 +92,6 @@ class ModelContextProtocolController < ActionController::API
       config.name = "MyMCPServer"
       config.title = "My MCP Server"
       config.version = "1.0.0"
-      config.logging_enabled = true
       config.registry = build_registry
       config.context = {
         user_id: current_user.id,
@@ -210,9 +210,6 @@ server = ModelContextProtocol::Server.new do |config|
     Use this server when you need to interact with the local development environment.
   INSTRUCTIONS
 
-  # Enable or disable MCP server logging
-  config.logging_enabled = true
-
   # Configure pagination options for the following methods:
   # prompts/list, resources/list, resource_template/list, tools/list
   config.pagination = {
@@ -249,11 +246,11 @@ server = ModelContextProtocol::Server.new do |config|
 
   # Register prompts, resources, resource templates, and tools
   config.registry = ModelContextProtocol::Server::Registry.new do
-    prompts list_changed: true do
+    prompts do
       register TestPrompt
     end
 
-    resources list_changed: true, subscribe: true do
+    resources do
       register TestResource
     end
 
@@ -261,7 +258,7 @@ server = ModelContextProtocol::Server.new do |config|
       register TestResourceTemplate
     end
 
-    tools list_changed: true do
+    tools do
       register TestTool
     end
   end
@@ -281,7 +278,6 @@ The following table details all available configuration options for the MCP serv
 | `version` | String | Yes | - | Version of the MCP server |
 | `title` | String | No | - | Human-readable display name for the MCP server |
 | `instructions` | String | No | - | Instructions for how the MCP server should be used by LLMs |
-| `logging_enabled` | Boolean | No | `true` | Enable or disable MCP server logging |
 | `pagination` | Hash/Boolean | No | See pagination table | Pagination configuration (or `false` to disable) |
 | `context` | Hash | No | `{}` | Contextual variables available to prompts, resources, and tools |
 | `transport` | Hash | No | `{ type: :stdio }` | Transport configuration |
@@ -347,6 +343,30 @@ end
 | `reaper_interval` | Integer | No | `60` | Reaper check interval in seconds |
 | `idle_timeout` | Integer | No | `300` | Idle connection timeout in seconds |
 
+### Server Logging Configuration
+
+Server logging can be configured globally to customize how your MCP server writes debug and operational logs. This logging is separate from client logging (which sends messages to MCP clients via the protocol) and is used for server-side debugging, monitoring, and troubleshooting:
+
+```ruby
+ModelContextProtocol::Server.configure_server_logging do |config|
+  config.logdev = $stderr           # or a file path like '/var/log/mcp-server.log'
+  config.level = Logger::INFO       # Logger::DEBUG, Logger::INFO, Logger::WARN, Logger::ERROR, Logger::FATAL
+  config.progname = "MyMCPServer"   # Program name for log entries
+  config.formatter = proc do |severity, datetime, progname, msg|
+    "#{datetime.strftime('%Y-%m-%d %H:%M:%S')} #{severity} [#{progname}] #{msg}\n"
+  end
+end
+```
+
+| Option | Type | Required | Default | Description |
+|--------|------|----------|---------|-------------|
+| `logdev` | IO/String | No | `$stderr` | Log destination (IO object or file path) |
+| `level` | Integer | No | `Logger::INFO` | Minimum log level to output |
+| `progname` | String | No | `"MCP-Server"` | Program name for log entries |
+| `formatter` | Proc | No | Default timestamp format | Custom log formatter |
+
+**Note:** When using `:stdio` transport, server logging must not use `$stdout` as it conflicts with the MCP protocol communication. Use `$stderr` or a file instead.
+
 ### Registry Configuration Options
 
 The registry is configured using `ModelContextProtocol::Server::Registry.new` and supports the following block types:
@@ -360,15 +380,17 @@ The registry is configured using `ModelContextProtocol::Server::Registry.new` an
 
 Within each block, use `register ClassName` to register your handlers.
 
+**Note:** The `list_changed` and `subscribe` options are accepted for capability advertisement but the  list changed notification functionality is not yet implemented (see [Feature Support](#feature-support-server)).
+
 **Example:**
 ```ruby
 config.registry = ModelContextProtocol::Server::Registry.new do
-  prompts list_changed: true do
+  prompts do
     register MyPrompt
     register AnotherPrompt
   end
 
-  resources list_changed: true, subscribe: true do
+  resources do
     register MyResource
   end
 
@@ -384,9 +406,9 @@ end
 
 The `ModelContextProtocol::Server::Prompt` base class allows subclasses to define a prompt that the MCP client can use.
 
-Define the prompt properties and then implement the `call` method to build your prompt. Any arguments passed to the tool from the MCP client will be available in the `arguments` hash with symbol keys (e.g., `arguments[:argument_name]`), and any context values provided in the server configuration will be available in the `context` hash. Use the `respond_with` instance method to ensure your prompt responds with appropriately formatted response data.
+Define the prompt properties and then implement the `call` method to build your prompt. Any arguments passed to the prompt from the MCP client will be available in the `arguments` hash with symbol keys (e.g., `arguments[:argument_name]`), and any context values provided in the server configuration will be available in the `context` hash. Use the `respond_with` instance method to ensure your prompt responds with appropriately formatted response data.
 
-You can also log from within your prompt by calling a valid logger level method on the `logger` and passing a string message.
+You can also send MCP log messages to clients from within your prompt by calling a valid logger level method on the `client_logger` and passing a string message. For server-side debugging and monitoring, use the `server_logger` to write logs that are not sent to clients.
 
 ### Prompt Definition
 
@@ -452,7 +474,8 @@ The `arguments` passed from an MCP client are available, as well as the `context
 |----------|---------|-------------|
 | `arguments` | Within `call` | Hash containing client-provided arguments (symbol keys) |
 | `context` | Within `call` | Hash containing server configuration context values |
-| `logger` | Within `call` | Logger instance for logging (e.g., `logger.info("message")`) |
+| `client_logger` | Within `call` | Client logger instance for sending MCP log messages (e.g., `client_logger.info("message")`) |
+| `server_logger` | Within `call` | Server logger instance for debugging and monitoring (e.g., `server_logger.debug("message")`) |
 
 ### Examples
 
@@ -505,8 +528,12 @@ class TestPrompt < ModelContextProtocol::Server::Prompt
 
   # The call method is invoked by the MCP Server to generate a response to resource/read requests
   def call
-    # You can use the logger
-    logger.info("Brainstorming excuses...")
+    # You can use the client_logger
+    client_logger.info("Brainstorming excuses...")
+
+    # Server logging for debugging and monitoring (not sent to client)
+    server_logger.debug("Prompt called with arguments: #{arguments}")
+    server_logger.info("Generating excuse brainstorming prompt")
 
     # Build an array of user and assistant messages
     messages = message_history do
@@ -541,6 +568,8 @@ The `ModelContextProtocol::Server::Resource` base class allows subclasses to def
 
 Define the resource properties and optionally annotations, then implement the `call` method to build your resource. Use the `respond_with` instance method to ensure your resource responds with appropriately formatted response data.
 
+You can also send MCP log messages to clients from within your resource by calling a valid logger level method on the `client_logger` and passing a string message. For server-side debugging and monitoring, use the `server_logger` to write logs that are not sent to clients.
+
 ### Resource Definition
 
 Use the `define` block to set [resource properties](https://spec.modelcontextprotocol.io/specification/2025-06-18/server/resources/) and configure annotations.
@@ -578,12 +607,15 @@ Define your resource properties and annotations, implement the `call` method to
 
 ### Available Instance Variables
 
-Resources are stateless and only have access to their configured properties.
+Resources have access to their configured properties and server context.
 
 | Variable | Context | Description |
 |----------|---------|-------------|
 | `mime_type` | Within `call` | The configured MIME type for this resource |
 | `uri` | Within `call` | The configured URI identifier for this resource |
+| `client_logger` | Within `call` | Client logger instance for sending MCP log messages (e.g., `client_logger.info("message")`) |
+| `server_logger` | Within `call` | Server logger instance for debugging and monitoring (e.g., `server_logger.debug("message")`) |
+| `context` | Within `call` | Hash containing server configuration context values |
 
 ### Examples
 
@@ -724,7 +756,9 @@ end
 
 The `ModelContextProtocol::Server::Tool` base class allows subclasses to define a tool that the MCP client can use.
 
-Define the tool properties and schemas, then implement the `call` method to build your tool response. Arguments from the MCP client and server context are available, along with logging capabilities.
+Define the tool properties and schemas, then implement the `call` method to build your tool. Any arguments passed to the tool from the MCP client will be available in the `arguments` hash with symbol keys (e.g., `arguments[:argument_name]`), and any context values provided in the server configuration will be available in the `context` hash. Use the `respond_with` instance method to ensure your prompt responds with appropriately formatted response data.
+
+You can also send MCP log messages to clients from within your tool by calling a valid logger level method on the `client_logger` and passing a string message. For server-side debugging and monitoring, use the `server_logger` to write logs that are not sent to clients.
 
 ### Tool Definition
 
@@ -781,7 +815,8 @@ Arguments from MCP clients and server context are available, along with logging
 |----------|---------|-------------|
 | `arguments` | Within `call` | Hash containing client-provided arguments (symbol keys) |
 | `context` | Within `call` | Hash containing server configuration context values |
-| `logger` | Within `call` | Logger instance for logging (e.g., `logger.info("message")`) |
+| `client_logger` | Within `call` | Client logger instance for sending MCP log messages (e.g., `client_logger.info("message")`) |
+| `server_logger` | Within `call` | Server logger instance for debugging and monitoring (e.g., `server_logger.debug("message")`) |
 
 ### Examples
 
@@ -835,11 +870,11 @@ class TestToolWithStructuredContentResponse < ModelContextProtocol::Server::Tool
   def call
     # Use values provided by the server as context
     user_id = context[:user_id]
-    logger.info("Initiating request for user #{user_id}...")
+    client_logger.info("Initiating request for user #{user_id}...")
 
     # Use values provided by clients as tool arguments
     location = arguments[:location]
-    logger.info("Getting weather data for #{location}...")
+    client_logger.info("Getting weather data for #{location}...")
 
     # Returns a hash that validates against the output schema
     weather_data = get_weather_data(location)
@@ -883,7 +918,7 @@ class TestToolWithTextResponse < ModelContextProtocol::Server::Tool
   end
 
   def call
-    logger.info("Silly user doesn't know how to double a number")
+    client_logger.info("Silly user doesn't know how to double a number")
     number = arguments[:number].to_i
     calculation = number * 2
 
@@ -975,7 +1010,7 @@ class TestToolWithResourceResponse < ModelContextProtocol::Server::Tool
       return respond_with :error, text: "Resource `#{name}` not found"
     end
 
-    resource_data = resource_klass.call
+    resource_data = resource_klass.call(client_logger, context)
 
     respond_with content: embedded_resource_content(resource: resource_data)
   end
@@ -1003,7 +1038,7 @@ class TestToolWithMixedContentResponse < ModelContextProtocol::Server::Tool
   end
 
   def call
-    logger.info("Getting comprehensive temperature history data")
+    client_logger.info("Getting comprehensive temperature history data")
 
     zip = arguments[:zip]
     temperature_history = retrieve_temperature_history(zip:)
@@ -1085,7 +1120,7 @@ class TestToolWithCancellableSleep < ModelContextProtocol::Server::Tool
   end
 
   def call
-    logger.info("Starting 3 second sleep operation")
+    client_logger.info("Starting 3 second sleep operation")
 
     result = cancellable do
       sleep 3
@@ -1126,7 +1161,7 @@ class TestToolWithProgressableAndCancellable < ModelContextProtocol::Server::Too
   def call
     max_duration = arguments[:max_duration] || 10
     work_steps = arguments[:work_steps] || 10
-    logger.info("Starting progressable call with max_duration=#{max_duration}, work_steps=#{work_steps}")
+    client_logger.info("Starting progressable call with max_duration=#{max_duration}, work_steps=#{work_steps}")
 
     result = progressable(max_duration:, message: "Processing #{work_steps} items") do
       cancellable do

data/lib/model_context_protocol/server/cancellable.rb

@@ -24,8 +24,8 @@ module ModelContextProtocol
       executing_thread = Concurrent::AtomicReference.new(nil)
 
       timer_task = Concurrent::TimerTask.new(execution_interval: interval) do
-        if context && context[:request_store] && context[:request_id]
-          if context[:request_store].cancelled?(context[:request_id])
+        if context && context[:request_store] && context[:jsonrpc_request_id]
+          if context[:request_store].cancelled?(context[:jsonrpc_request_id])
             thread = executing_thread.get
             thread&.raise(CancellationError, "Request was cancelled") if thread&.alive?
           end
@@ -35,9 +35,9 @@ module ModelContextProtocol
       begin
         executing_thread.set(Thread.current)
 
-        if context && context[:request_store] && context[:request_id]
-          if context[:request_store].cancelled?(context[:request_id])
-            raise CancellationError, "Request #{context[:request_id]} was cancelled"
+        if context && context[:request_store] && context[:jsonrpc_request_id]
+          if context[:request_store].cancelled?(context[:jsonrpc_request_id])
+            raise CancellationError, "Request #{context[:jsonrpc_request_id]} was cancelled"
           end
         end
 

data/lib/model_context_protocol/server/{mcp_logger.rb → client_logger.rb}

@@ -3,11 +3,13 @@ require "forwardable"
 require "json"
 
 module ModelContextProtocol
-  class Server::MCPLogger
+  class Server::ClientLogger
     extend Forwardable
 
     def_delegators :@internal_logger, :datetime_format=, :formatter=, :progname, :progname=
 
+    VALID_LOG_LEVELS = %w[debug info notice warning error critical alert emergency].freeze
+
     LEVEL_MAP = {
       "debug" => Logger::DEBUG,
       "info" => Logger::INFO,
@@ -29,11 +31,10 @@ module ModelContextProtocol
     }.freeze
 
     attr_accessor :transport
-    attr_reader :logger_name, :enabled
+    attr_reader :logger_name
 
-    def initialize(logger_name: "server", level: "info", enabled: true)
+    def initialize(logger_name: "server", level: "info")
       @logger_name = logger_name
-      @enabled = enabled
       @internal_logger = Logger.new(nil)
       @internal_logger.level = LEVEL_MAP[level] || Logger::INFO
       @transport = nil
@@ -42,13 +43,11 @@ module ModelContextProtocol
 
     %i[debug info warn error fatal unknown].each do |severity|
       define_method(severity) do |message = nil, **data, &block|
-        return true unless @enabled
         add(Logger.const_get(severity.to_s.upcase), message, data, &block)
       end
     end
 
     def add(severity, message = nil, data = {}, &block)
-      return true unless @enabled
       return true if severity < @internal_logger.level
 
       message = block.call if message.nil? && block_given?
@@ -70,14 +69,12 @@ module ModelContextProtocol
 
     def connect_transport(transport)
       @transport = transport
-      flush_queued_messages if @enabled
+      flush_queued_messages
     end
 
     private
 
     def send_notification(severity, message, data)
-      return unless @enabled
-
       notification_params = {
         level: REVERSE_LEVEL_MAP[severity] || "info",
         logger: @logger_name,
@@ -99,7 +96,7 @@ module ModelContextProtocol
     end
 
     def flush_queued_messages
-      return unless @transport && @enabled
+      return unless @transport
       @queued_messages.each do |params|
         @transport.send_notification("notifications/message", params)
       end

data/lib/model_context_protocol/server/configuration.rb

@@ -1,5 +1,3 @@
-require_relative "mcp_logger"
-
 module ModelContextProtocol
   class Server::Configuration
     # Raised when configured with invalid name.
@@ -23,48 +21,25 @@ module ModelContextProtocol
     # Raised when transport configuration is invalid
     class InvalidTransportError < StandardError; end
 
-    # Raised when an invalid log level is provided
-    class InvalidLogLevelError < StandardError; end
-
     # Raised when pagination configuration is invalid
     class InvalidPaginationError < StandardError; end
 
-    # Valid MCP log levels per the specification
-    VALID_LOG_LEVELS = %w[debug info notice warning error critical alert emergency].freeze
-
     attr_accessor :name, :registry, :version, :transport, :pagination, :title, :instructions
-    attr_reader :logger
+    attr_reader :client_logger, :server_logger
 
     def initialize
-      @logging_enabled = true
-      @default_log_level = "info"
-      @logger = ModelContextProtocol::Server::MCPLogger.new(
+      @client_logger = ModelContextProtocol::Server::ClientLogger.new(
         logger_name: "server",
-        level: @default_log_level,
-        enabled: @logging_enabled
+        level: "info"
       )
-    end
 
-    def logging_enabled?
-      @logging_enabled
-    end
-
-    def logging_enabled=(value)
-      @logging_enabled = value
-      @logger = ModelContextProtocol::Server::MCPLogger.new(
-        logger_name: "server",
-        level: @default_log_level,
-        enabled: value
-      )
-    end
-
-    def default_log_level=(level)
-      unless VALID_LOG_LEVELS.include?(level.to_s)
-        raise InvalidLogLevelError, "Invalid log level: #{level}. Valid levels are: #{VALID_LOG_LEVELS.join(", ")}"
+      server_logger_params = if ModelContextProtocol::Server::GlobalConfig::ServerLogging.configured?
+        ModelContextProtocol::Server::GlobalConfig::ServerLogging.logger_params
+      else
+        {}
       end
 
-      @default_log_level = level.to_s
-      @logger.set_mcp_level(@default_log_level)
+      @server_logger = ModelContextProtocol::Server::ServerLogger.new(**server_logger_params)
     end
 
     def transport_type
@@ -123,12 +98,13 @@ module ModelContextProtocol
       raise InvalidServerNameError unless valid_name?
       raise InvalidRegistryError unless valid_registry?
       raise InvalidServerVersionError unless valid_version?
+
       validate_transport!
       validate_pagination!
       validate_title!
       validate_instructions!
-
       validate_environment_variables!
+      validate_server_logging_transport_constraints!
     end
 
     def environment_variables
@@ -176,6 +152,13 @@ module ModelContextProtocol
       end
     end
 
+    def validate_server_logging_transport_constraints!
+      return unless transport_type == :stdio && server_logger.logdev == $stdout
+
+      raise ModelContextProtocol::Server::ServerLogger::StdoutNotAllowedError,
+        "StdioTransport cannot log to stdout. Use stderr or a file instead."
+    end
+
     def valid_name?
       name&.is_a?(String)
     end

data/lib/model_context_protocol/server/global_config/server_logging.rb

@@ -0,0 +1,78 @@
+require "singleton"
+
+module ModelContextProtocol
+  module Server::GlobalConfig
+    class ServerLogging
+      include Singleton
+
+      class NotConfiguredError < StandardError
+        def initialize
+          super("Server logging not configured. Call ModelContextProtocol::Server.configure_server_logging first")
+        end
+      end
+
+      class LoggerConfig
+        attr_accessor :logdev, :level, :formatter, :progname
+
+        def initialize
+          @level = Logger::INFO
+          @progname = "MCP-Server"
+        end
+
+        def to_h
+          {
+            logdev: @logdev,
+            level: @level,
+            formatter: @formatter,
+            progname: @progname
+          }.compact
+        end
+      end
+
+      def self.configure(&block)
+        instance.configure(&block)
+      end
+
+      def self.configured?
+        instance.configured?
+      end
+
+      def self.logger_params
+        instance.logger_params
+      end
+
+      def self.reset!
+        instance.reset!
+      end
+
+      def configure(&block)
+        raise ArgumentError, "Configuration block required" unless block_given?
+
+        @config = LoggerConfig.new
+        yield(@config)
+        @configured = true
+      end
+
+      def configured?
+        @configured == true
+      end
+
+      def logger_params
+        raise NotConfiguredError unless configured?
+
+        @config.to_h
+      end
+
+      def reset!
+        @configured = false
+        @config = nil
+      end
+
+      private
+
+      def initialize
+        reset!
+      end
+    end
+  end
+end

data/lib/model_context_protocol/server/progressable.rb

@@ -16,56 +16,78 @@ module ModelContextProtocol
     #   end
     def progressable(max_duration:, message: nil, &block)
       context = Thread.current[:mcp_context]
-
       return yield unless context && context[:progress_token] && context[:transport]
 
       progress_token = context[:progress_token]
       transport = context[:transport]
+      jsonrpc_request_id = context[:jsonrpc_request_id]
+      request_store = context[:request_store]
+      stream_id = context[:stream_id]
+
       start_time = Time.now
       update_interval = [1.0, max_duration * 0.05].max
 
       timer_task = Concurrent::TimerTask.new(execution_interval: update_interval) do
-        elapsed_seconds = Time.now - start_time
-        progress_pct = [(elapsed_seconds / max_duration) * 100, 99].min
-
-        progress_message = if message
-          "#{message} (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
-        else
-          "Processing... (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
-        end
+        Thread.current[:mcp_context] = {jsonrpc_request_id:}
 
         begin
-          transport.send_notification("notifications/progress", {
-            progressToken: progress_token,
-            progress: progress_pct.round(1),
-            total: 100,
-            message: progress_message
-          })
-        rescue
-          nil
-        end
+          if request_store && jsonrpc_request_id
+            break if request_store.cancelled?(jsonrpc_request_id)
+          end
+
+          elapsed_seconds = Time.now - start_time
+          progress_pct = [(elapsed_seconds / max_duration) * 100, 99].min
+
+          progress_message = if message
+            "#{message} (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
+          else
+            "Processing... (#{elapsed_seconds.round(1)}s / ~#{max_duration}s)"
+          end
+
+          begin
+            transport.send_notification("notifications/progress", {
+              progressToken: progress_token,
+              progress: progress_pct.round(1),
+              total: 100,
+              message: progress_message
+            }, session_id: stream_id)
+          rescue
+            break
+          end
 
-        timer_task.shutdown if elapsed_seconds >= max_duration
+          timer_task.shutdown if elapsed_seconds >= max_duration
+        ensure
+          Thread.current[:mcp_context] = nil
+        end
       end
 
       begin
         timer_task.execute
+
         result = yield
 
+        original_context = Thread.current[:mcp_context]
+        Thread.current[:mcp_context] = {jsonrpc_request_id:}
+
         begin
           transport.send_notification("notifications/progress", {
             progressToken: progress_token,
             progress: 100,
             total: 100,
             message: "Completed"
-          })
+          }, session_id: stream_id)
         rescue
           nil
+        ensure
+          Thread.current[:mcp_context] = original_context
         end
 
         result
       ensure
-        timer_task&.shutdown if timer_task&.running?
+        if timer_task&.running?
+          timer_task.shutdown
+          sleep(0.1) if timer_task.running?
+        end
       end
     end
   end