lib/angelo/base.rb in angelo-0.2.1 vs lib/angelo/base.rb in angelo-0.2.2

@@ -126,19 +126,19 @@
 
       def task name, &block
         Angelo::Server.define_task name, &block
       end
 
-      def websockets
+      def websockets reject = true
         @websockets ||= Stash::Websocket.new server
-        @websockets.reject! &:closed?
+        @websockets.reject! &:closed? if reject
         @websockets
       end
 
-      def sses
+      def sses reject = true
         @sses ||= Stash::SSE.new server
-        @sses.reject! &:closed?
+        @sses.reject! &:closed? if reject
         @sses
       end
 
       def content_type type
         Responder.content_type type
@@ -285,11 +285,11 @@
       halt 200, data
     end
 
     def eventsource &block
       headers SSE_HEADER
-      async :handle_event_source, EventSource.new(responder.connection.detach.socket), block
+      async :handle_event_source, EventSource.new(responder), block
       halt 200, :sse
     end
 
     def report_errors?
       @@report_errors
@@ -311,22 +311,32 @@
     end
 
     class EventSource
       extend Forwardable
 
-      def_delegators :@socket, :close, :closed?, :<<, :write
-      attr_reader :socket
+      def_delegators :@socket, :close, :closed?, :<<, :write, :peeraddr
+      attr_reader :responder, :socket
 
-      def initialize socket
-        @socket = socket
+      def initialize responder
+        @responder = responder
+        @socket = @responder.connection.detach.socket
       end
 
       def event name, data = nil
         @socket.write Base.sse_event(name, data)
       end
 
       def message data
         @socket.write Base.sse_message(data)
+      end
+
+      def on_close &block
+        @responder.on_close = block
+      end
+
+      def on_close= block
+        raise ArgumentError.new unless Proc === block
+        @responder.on_close = block
       end
 
     end
 
     class ChunkedResponse