Class: ModelContextProtocol::Server::StreamableHttpTransport::RequestStore

Inherits:
Object
  • Object
show all
Defined in:
lib/model_context_protocol/server/streamable_http_transport/request_store.rb

Overview

Redis-based distributed storage for tracking active requests and their cancellation status. This store is used by StreamableHttpTransport to manage request lifecycle across multiple server instances and handle cancellation in a distributed environment.

Constant Summary collapse

REQUEST_KEY_PREFIX = 
"request:active:"
CANCELLED_KEY_PREFIX = 
"request:cancelled:"
SESSION_KEY_PREFIX = 
"request:session:"
DEFAULT_TTL =

1 minute TTL for request entries

60

Instance Method Summary collapse

Constructor Details

#initialize(redis_client, server_instance, ttl: DEFAULT_TTL) ⇒ RequestStore

Returns a new instance of RequestStore.



14
15
16
17
18
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 14

def initialize(redis_client, server_instance, ttl: DEFAULT_TTL)
  @redis = redis_client
  @server_instance = server_instance
  @ttl = ttl
end

Instance Method Details

#cancelled?(jsonrpc_request_id) ⇒ Boolean

Check if a request has been cancelled

Parameters:

  • jsonrpc_request_id (String) —

    the unique JSON-RPC request identifier

Returns:

  • (Boolean) —

    true if the request is cancelled, false otherwise



63
64
65
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 63

def cancelled?(jsonrpc_request_id)
  @redis.exists("#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}") == 1
end

#cleanup_session_requests(session_id) ⇒ Array<String>

Clean up all requests associated with a session This is typically called when a session is terminated

Parameters:

  • session_id (String) —

    the session identifier

Returns:

  • (Array<String>) —

    list of cleaned up request IDs



98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 98

def cleanup_session_requests(session_id)
  pattern = "#{SESSION_KEY_PREFIX}#{session_id}:*"
  request_keys = @redis.keys(pattern)
  return [] if request_keys.empty?

  # Extract request IDs from the keys
  jsonrpc_request_ids = request_keys.map do |key|
    key.sub("#{SESSION_KEY_PREFIX}#{session_id}:", "")
  end

  # Delete all related keys
  all_keys = []
  jsonrpc_request_ids.each do |jsonrpc_request_id|
    all_keys << "#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}"
    all_keys << "#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}"
  end
  all_keys.concat(request_keys)

  @redis.del(*all_keys) unless all_keys.empty?
  jsonrpc_request_ids
end

#mark_cancelled(jsonrpc_request_id, reason = nil) ⇒ Boolean

Mark a request as cancelled

Parameters:

  • jsonrpc_request_id (String) —

    the unique JSON-RPC request identifier

  • reason (String) (defaults to: nil) —

    optional reason for cancellation

Returns:

  • (Boolean) —

    true if cancellation was recorded



48
49
50
51
52
53
54
55
56
57
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 48

def mark_cancelled(jsonrpc_request_id, reason = nil)
  cancellation_data = {
    cancelled_at: Time.now.to_f,
    reason: reason
  }

  result = @redis.set("#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}",
    cancellation_data.to_json, ex: @ttl)
  result == "OK"
end

#register_request(jsonrpc_request_id, session_id = nil) ⇒ void

This method returns an undefined value.

Register a new request with its associated session

Parameters:

  • jsonrpc_request_id (String) —

    the unique JSON-RPC request identifier

  • session_id (String) (defaults to: nil) —

    the session identifier (can be nil for sessionless requests)



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 25

def register_request(jsonrpc_request_id, session_id = nil)
  request_data = {
    session_id: session_id,
    server_instance: @server_instance,
    started_at: Time.now.to_f
  }

  @redis.multi do |multi|
    multi.set("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}",
      request_data.to_json, ex: @ttl)

    if session_id
      multi.set("#{SESSION_KEY_PREFIX}#{session_id}:#{jsonrpc_request_id}",
        true, ex: @ttl)
    end
  end
end

#unregister_request(jsonrpc_request_id) ⇒ void

This method returns an undefined value.

Unregister a request (typically called when request completes)

Parameters:

  • jsonrpc_request_id (String) —

    the unique JSON-RPC request identifier



71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
 
# File 'lib/model_context_protocol/server/streamable_http_transport/request_store.rb', line 71

def unregister_request(jsonrpc_request_id)
  request_data = @redis.get("#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}")

  keys_to_delete = ["#{REQUEST_KEY_PREFIX}#{jsonrpc_request_id}",
    "#{CANCELLED_KEY_PREFIX}#{jsonrpc_request_id}"]

  if request_data
    begin
      data = JSON.parse(request_data)
      session_id = data["session_id"]

      if session_id
        keys_to_delete << "#{SESSION_KEY_PREFIX}#{session_id}:#{jsonrpc_request_id}"
      end
    rescue JSON::ParserError
      nil
    end
  end

  @redis.del(*keys_to_delete) unless keys_to_delete.empty?
end