Fix keepAlive reconnection (dart-lang/sse#23)

diff --git a/pkgs/sse/CHANGELOG.md b/pkgs/sse/CHANGELOG.md
index a318705..5135fb4 100644
--- a/pkgs/sse/CHANGELOG.md
+++ b/pkgs/sse/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 3.2.1
+
+- Fix an issue where `keepAlive` would only allow a single reconnection.
+
 ## 3.2.0
 
 - Re-expose `isInKeepAlivePeriod` flag on `SseConnection`. This flag will be

diff --git a/pkgs/sse/lib/src/server/sse_handler.dart b/pkgs/sse/lib/src/server/sse_handler.dart
index e8cd523..1749f46 100644
--- a/pkgs/sse/lib/src/server/sse_handler.dart
+++ b/pkgs/sse/lib/src/server/sse_handler.dart

@@ -163,16 +163,15 @@
         unawaited(connection._closedCompleter.future.then((_) {
           _connections.remove(clientId);
         }));
-        // Remove connection when it is remotely closed or the stream is
-        // cancelled.
-        channel.stream.listen((_) {
-          // SSE is unidirectional. Responses are handled through POST requests.
-        }, onDone: () {
-          connection._handleDisconnect();
-        });
-
         _connectionController.add(connection);
       }
+      // Remove connection when it is remotely closed or the stream is
+      // cancelled.
+      channel.stream.listen((_) {
+        // SSE is unidirectional. Responses are handled through POST requests.
+      }, onDone: () {
+        _connections[clientId]?._handleDisconnect();
+      });
     });
     return shelf.Response.notFound('');
   }

diff --git a/pkgs/sse/pubspec.yaml b/pkgs/sse/pubspec.yaml
index 9e8225a..1381d61 100644
--- a/pkgs/sse/pubspec.yaml
+++ b/pkgs/sse/pubspec.yaml

@@ -1,5 +1,5 @@
 name: sse
-version: 3.2.0
+version: 3.2.1
 homepage: https://github.com/dart-lang/sse
 description: >-
   Provides client and server functionality for setting up bi-directional

diff --git a/pkgs/sse/test/sse_test.dart b/pkgs/sse/test/sse_test.dart
index 4f66393..3a7f76b 100644
--- a/pkgs/sse/test/sse_test.dart
+++ b/pkgs/sse/test/sse_test.dart

@@ -6,6 +6,7 @@
 import 'dart:async';
 import 'dart:io';
 
+import 'package:async/async.dart';
 import 'package:shelf/shelf.dart' as shelf;
 import 'package:shelf/shelf_io.dart' as io;
 import 'package:shelf_static/shelf_static.dart';
@@ -195,7 +196,20 @@
       // Ensure we can still round-trip data on the original connection and that
       // the connection is no longer marked keep-alive once it's reconnected.
       connection.sink.add('bar');
-      expect(await connection.stream.first, 'bar');
+      var queue = StreamQueue(connection.stream);
+      expect(await queue.next, 'bar');
+      expect(connection.isInKeepAlivePeriod, isFalse);
+
+      // Now check that we can reconnect multiple times.
+      closeSink(connection);
+      maxPumps = 50;
+      while (!connection.isInKeepAlivePeriod && maxPumps-- > 0) {
+        await pumpEventQueue(times: 1);
+      }
+      expect(connection.isInKeepAlivePeriod, isTrue);
+      expect(handler.numberOfClients, 1);
+      connection.sink.add('bar');
+      expect(await queue.next, 'bar');
       expect(connection.isInKeepAlivePeriod, isFalse);
     });