Class: ModelContextProtocol::Server::StreamableHttpTransport::StreamRegistry

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

Constant Summary collapse

STREAM_KEY_PREFIX = 
"stream:active:"
HEARTBEAT_KEY_PREFIX = 
"stream:heartbeat:"
DEFAULT_TTL =

1 minute TTL for stream entries

60

Instance Method Summary collapse

Constructor Details

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

Returns a new instance of StreamRegistry.



12
13
14
15
16
17
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 12

def initialize(redis_client, server_instance, ttl: DEFAULT_TTL)
  @redis = redis_client
  @server_instance = server_instance
  @ttl = ttl
  @local_streams = {} # Keep local reference for direct stream access
end

Instance Method Details

#cleanup_expired_streams ⇒ Object 



61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 61

def cleanup_expired_streams
  # Get all local stream session IDs
  local_session_ids = @local_streams.keys

  # Check which ones are still active in Redis
  pipeline_results = @redis.pipelined do |pipeline|
    local_session_ids.each do |session_id|
      pipeline.exists("#{STREAM_KEY_PREFIX}#{session_id}")
    end
  end

  # Remove expired streams from local storage
  expired_sessions = []
  local_session_ids.each_with_index do |session_id, index|
    if pipeline_results[index] == 0 # Stream expired in Redis
      @local_streams.delete(session_id)
      expired_sessions << session_id
    end
  end

  expired_sessions
end

#get_all_local_streams ⇒ Object 



53
54
55
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 53

def get_all_local_streams
  @local_streams.dup
end

#get_local_stream(session_id) ⇒ Object 



38
39
40
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 38

def get_local_stream(session_id)
  @local_streams[session_id]
end

#has_any_local_streams? ⇒ Boolean

Returns:



57
58
59
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 57

def has_any_local_streams?
  !@local_streams.empty?
end

#has_local_stream?(session_id) ⇒ Boolean

Returns:



42
43
44
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 42

def has_local_stream?(session_id)
  @local_streams.key?(session_id)
end

#refresh_heartbeat(session_id) ⇒ Object 



46
47
48
49
50
51
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 46

def refresh_heartbeat(session_id)
  @redis.multi do |multi|
    multi.set("#{HEARTBEAT_KEY_PREFIX}#{session_id}", Time.now.to_f, ex: @ttl)
    multi.expire("#{STREAM_KEY_PREFIX}#{session_id}", @ttl)
  end
end

#register_stream(session_id, stream) ⇒ Object 



19
20
21
22
23
24
25
26
27
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 19

def register_stream(session_id, stream)
  @local_streams[session_id] = stream

  # Store stream registration in Redis with TTL
  @redis.multi do |multi|
    multi.set("#{STREAM_KEY_PREFIX}#{session_id}", @server_instance, ex: @ttl)
    multi.set("#{HEARTBEAT_KEY_PREFIX}#{session_id}", Time.now.to_f, ex: @ttl)
  end
end

#unregister_stream(session_id) ⇒ Object 



29
30
31
32
33
34
35
36
 
# File 'lib/model_context_protocol/server/streamable_http_transport/stream_registry.rb', line 29

def unregister_stream(session_id)
  @local_streams.delete(session_id)

  @redis.multi do |multi|
    multi.del("#{STREAM_KEY_PREFIX}#{session_id}")
    multi.del("#{HEARTBEAT_KEY_PREFIX}#{session_id}")
  end
end