feat(sse): implement Sinatra SSE + improvement on connection handling

Author: Dogan AY

packages/forest_admin_datasource_rpc/lib/forest_admin_datasource_rpc/Utils/sse_client.rb

@@ -6,12 +6,15 @@
 module ForestAdminDatasourceRpc
   module Utils
     class SseClient
+      attr_reader :closed
+
       def initialize(uri, auth_secret, &on_rpc_stop)
         @uri = uri
         @auth_secret = auth_secret
         @on_rpc_stop = on_rpc_stop
         @client = nil
         @closed = false
+        @connection_attempts = 0
       end
 
       def start
@@ -26,26 +29,47 @@ def start
           'X_SIGNATURE' => signature
         }
 
-        ForestAdminRpcAgent::Facades::Container.logger.log('Debug', "Connecting to SSE at #{@uri}.")
+        @connection_attempts += 1
+        ForestAdminRpcAgent::Facades::Container.logger&.log(
+          'Debug',
+          "[SSE Client] Connecting to #{@uri} (attempt ##{@connection_attempts})"
+        )
 
-        @client = SSE::Client.new(@uri, headers: headers) do |client|
-          client.on_event do |event|
-            handle_event(event)
-          end
+        begin
+          @client = SSE::Client.new(@uri, headers: headers) do |client|
+            client.on_event do |event|
+              handle_event(event)
+            end
 
-          client.on_error do |err|
-            # TODO: optimisation on client close
-            # ForestAdminRpcAgent::Facades::Container.logger.log('Warn', "[SSE] Error: #{err.class} - #{err.message}")
+            client.on_error do |err|
+              handle_error(err)
+            end
           end
+
+          ForestAdminRpcAgent::Facades::Container.logger&.log('Debug', '[SSE Client] Connected successfully')
+        rescue StandardError => e
+          ForestAdminRpcAgent::Facades::Container.logger&.log(
+            'Error',
+            "[SSE Client] Failed to connect: #{e.class} - #{e.message}"
+          )
+          raise
         end
       end
 
       def close
         return if @closed
 
         @closed = true
-        @client&.close
-        # ForestAdminRpcAgent::Facades::Container.logger.log('Debug', '[SSE] Client closed')
+        ForestAdminRpcAgent::Facades::Container.logger&.log('Debug', '[SSE Client] Closing connection')
+
+        begin
+          @client&.close
+        rescue StandardError => e
+          ForestAdminRpcAgent::Facades::Container.logger&.log('Debug',
+                                                              "[SSE Client] Error during close: #{e.message}")
+        end
+
+        ForestAdminRpcAgent::Facades::Container.logger&.log('Debug', '[SSE Client] Connection closed')
       end
 
       private
@@ -56,14 +80,46 @@ def handle_event(event)
 
         case type
         when 'heartbeat'
-          # ForestAdminRpcAgent::Facades::Container.logger.log('Debug', '[SSE] Heartbeat')
+          # Heartbeat received - connection is alive
         when 'RpcServerStop'
-          # ForestAdminRpcAgent::Facades::Container.logger.log('Debug', '[SSE] RpcServerStop received')
-          @on_rpc_stop&.call
+          ForestAdminRpcAgent::Facades::Container.logger&.log('Debug', '[SSE Client] RpcServerStop received')
+          handle_rpc_stop
         else
-          ForestAdminRpcAgent::Facades::Container.logger.log('Debug',
-                                                             "[SSE] Unknown event: #{type} with payload: #{data}")
+          ForestAdminRpcAgent::Facades::Container.logger&.log(
+            'Debug',
+            "[SSE Client] Unknown event: #{type} with payload: #{data}"
+          )
         end
+      rescue StandardError => e
+        ForestAdminRpcAgent::Facades::Container.logger&.log(
+          'Error',
+          "[SSE Client] Error handling event: #{e.class} - #{e.message}"
+        )
+      end
+
+      def handle_error(err)
+        # Ignore errors when client is intentionally closed
+        return if @closed
+
+        error_message = case err
+                        when EOFError, IOError
+                          "Connection lost: #{err.class}"
+                        when StandardError
+                          "#{err.class} - #{err.message}"
+                        else
+                          err.to_s
+                        end
+
+        ForestAdminRpcAgent::Facades::Container.logger&.log('Warn', "[SSE Client] Error: #{error_message}")
+      end
+
+      def handle_rpc_stop
+        @on_rpc_stop&.call
+      rescue StandardError => e
+        ForestAdminRpcAgent::Facades::Container.logger&.log(
+          'Error',
+          "[SSE Client] Error in RPC stop callback: #{e.class} - #{e.message}"
+        )
       end
 
       def generate_signature(timestamp)

packages/forest_admin_datasource_rpc/spec/lib/forest_admin_datasource_rpc/utils/sse_client_spec.rb

@@ -12,6 +12,27 @@ module Utils
         allow(ForestAdminRpcAgent::Facades::Container).to receive(:logger).and_return(logger)
       end
 
+      describe '#initialize' do
+        it 'initializes with correct attributes' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          expect(client.instance_variable_get(:@uri)).to eq(uri)
+          expect(client.instance_variable_get(:@auth_secret)).to eq(secret)
+          expect(client.instance_variable_get(:@closed)).to be false
+          expect(client.instance_variable_get(:@connection_attempts)).to eq(0)
+        end
+
+        it 'exposes closed status via attr_reader' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          expect(client.closed).to be false
+
+          client.close
+
+          expect(client.closed).to be true
+        end
+      end
+
       describe '#start' do
         it 'connects to SSE with the expected headers' do
           fake_client = instance_double(SSE::Client)
@@ -49,6 +70,63 @@ module Utils
           # finally, we check that it didn't try to connect
           expect(SSE::Client).not_to have_received(:new)
         end
+
+        it 'increments connection attempts counter' do
+          fake_client = instance_double(SSE::Client)
+          allow(SSE::Client).to receive(:new).and_yield(fake_client).and_return(fake_client)
+          allow(fake_client).to receive(:on_event)
+          allow(fake_client).to receive(:on_error)
+
+          client = described_class.new(uri, secret) { callback.call }
+
+          expect(client.instance_variable_get(:@connection_attempts)).to eq(0)
+
+          client.start
+
+          expect(client.instance_variable_get(:@connection_attempts)).to eq(1)
+
+          # Simulate reconnection
+          client.instance_variable_set(:@closed, false)
+          client.start
+
+          expect(client.instance_variable_get(:@connection_attempts)).to eq(2)
+        end
+
+        it 'logs connection attempt with attempt number' do
+          fake_client = instance_double(SSE::Client)
+          allow(SSE::Client).to receive(:new).and_yield(fake_client)
+          allow(fake_client).to receive(:on_event)
+          allow(fake_client).to receive(:on_error)
+
+          client = described_class.new(uri, secret) { callback.call }
+          client.start
+
+          expect(logger).to have_received(:log).with('Debug', /Connecting to.*attempt #1/)
+        end
+
+        it 'logs successful connection' do
+          fake_client = instance_double(SSE::Client)
+          allow(SSE::Client).to receive(:new).and_yield(fake_client)
+          allow(fake_client).to receive(:on_event)
+          allow(fake_client).to receive(:on_error)
+
+          client = described_class.new(uri, secret) { callback.call }
+          client.start
+
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] Connected successfully')
+        end
+
+        it 'logs and re-raises connection errors' do
+          allow(SSE::Client).to receive(:new).and_raise(StandardError, 'Connection failed')
+
+          client = described_class.new(uri, secret) { callback.call }
+
+          expect do
+            client.start
+          end.to raise_error(StandardError, 'Connection failed')
+
+          expect(logger).to have_received(:log).with('Error', /Failed to connect.*StandardError/)
+        end
       end
 
       describe '#close' do
@@ -71,6 +149,45 @@ module Utils
 
           expect { client.close }.not_to raise_error
         end
+
+        it 'logs closing and closed messages' do
+          fake_client = instance_double(SSE::Client, close: true)
+          allow(SSE::Client).to receive(:new).and_return(fake_client)
+          allow(fake_client).to receive(:on_event)
+          allow(fake_client).to receive(:on_error)
+
+          client = described_class.new(uri, secret) { callback.call }
+          client.start
+          client.close
+
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] Closing connection')
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] Connection closed')
+        end
+
+        it 'handles errors during close gracefully' do
+          fake_client = instance_double(SSE::Client)
+          allow(SSE::Client).to receive(:new).and_return(fake_client)
+          allow(fake_client).to receive(:on_event)
+          allow(fake_client).to receive(:on_error)
+          allow(fake_client).to receive(:close).and_raise(StandardError, 'Close failed')
+
+          client = described_class.new(uri, secret) { callback.call }
+          client.start
+
+          expect { client.close }.not_to raise_error
+
+          expect(logger).to have_received(:log).with('Debug', /Error during close/)
+        end
+
+        it 'sets closed flag to true' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          expect(client.closed).to be false
+
+          client.close
+
+          expect(client.closed).to be true
+        end
       end
 
       describe '#handle_event' do
@@ -99,7 +216,112 @@ module Utils
           event = Struct.new(:type, :data).new('FooEvent', 'hello')
           client.send(:handle_event, event)
 
-          expect(logger).to have_received(:log).with('Debug', '[SSE] Unknown event: FooEvent with payload: hello')
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] Unknown event: FooEvent with payload: hello')
+        end
+
+        it 'logs RpcServerStop event when received' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          event = Struct.new(:type, :data).new('RpcServerStop', '')
+          client.send(:handle_event, event)
+
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] RpcServerStop received')
+        end
+
+        it 'handles errors during event processing' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          event = Struct.new(:type, :data).new('invalid', nil)
+          allow(event).to receive(:type).and_raise(StandardError, 'Event parsing failed')
+
+          expect { client.send(:handle_event, event) }.not_to raise_error
+
+          expect(logger).to have_received(:log).with('Error', /Error handling event/)
+        end
+
+        it 'strips whitespace from event type and data' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          event = Struct.new(:type, :data).new('  FooEvent  ', '  data  ')
+          client.send(:handle_event, event)
+
+          expect(logger).to have_received(:log).with('Debug', '[SSE Client] Unknown event: FooEvent with payload: data')
+        end
+      end
+
+      describe '#handle_error' do
+        it 'logs connection errors' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          error = StandardError.new('Connection interrupted')
+          client.send(:handle_error, error)
+
+          expect(logger).to have_received(:log).with('Warn', /Error: StandardError - Connection interrupted/)
+        end
+
+        it 'identifies EOFError as connection lost' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          error = EOFError.new
+          client.send(:handle_error, error)
+
+          expect(logger).to have_received(:log).with('Warn', /Connection lost: EOFError/)
+        end
+
+        it 'identifies IOError as connection lost' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          error = IOError.new
+          client.send(:handle_error, error)
+
+          expect(logger).to have_received(:log).with('Warn', /Connection lost: IOError/)
+        end
+
+        it 'does not log errors when client is closed' do
+          client = described_class.new(uri, secret) { callback.call }
+          client.close
+
+          error = StandardError.new('Error after close')
+          client.send(:handle_error, error)
+
+          expect(logger).not_to have_received(:log).with('Warn', anything)
+        end
+      end
+
+      describe '#handle_rpc_stop' do
+        it 'executes the callback safely' do
+          client = described_class.new(uri, secret) { callback.call }
+
+          client.send(:handle_rpc_stop)
+
+          expect(callback).to have_received(:call)
+        end
+
+        it 'handles callback errors without crashing' do
+          error_callback = proc { raise StandardError, 'Callback failed' }
+          client = described_class.new(uri, secret, &error_callback)
+
+          expect { client.send(:handle_rpc_stop) }.not_to raise_error
+
+          expect(logger).to have_received(:log).with('Error', /Error in RPC stop callback/)
+        end
+
+        it 'handles nil callback gracefully' do
+          client = described_class.new(uri, secret)
+
+          expect { client.send(:handle_rpc_stop) }.not_to raise_error
+        end
+      end
+
+      describe '#generate_signature' do
+        it 'generates correct HMAC signature' do
+          client = described_class.new(uri, secret) { callback.call }
+          timestamp = '2025-01-01T12:00:00Z'
+
+          signature = client.send(:generate_signature, timestamp)
+
+          expected_signature = OpenSSL::HMAC.hexdigest('SHA256', secret, timestamp)
+          expect(signature).to eq(expected_signature)
         end
       end
     end