Handle response for closed MVC WebSocket session

spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java

@@ -345,12 +345,6 @@ private GraphQlWebSocketMessage decode(TextMessage message) throws IOException {
 				.read(GraphQlWebSocketMessage.class, new HttpInputMessageAdapter(message));
 	}
 
-	private SessionState getSessionInfo(WebSocketSession session) {
-		SessionState info = this.sessionInfoMap.get(session.getId());
-		Assert.notNull(info, "No SessionInfo for " + session);
-		return info;
-	}
-
 	@SuppressWarnings("unchecked")
 	private Flux<TextMessage> handleResponse(WebSocketSession session, String id, WebGraphQlResponse response) {
 		if (logger.isDebugEnabled()) {
@@ -364,10 +358,15 @@ private Flux<TextMessage> handleResponse(WebSocketSession session, String id, We
 			responseFlux = Flux.from((Publisher<ExecutionResult>) response.getData())
 					.map(ExecutionResult::toSpecification)
 					.doOnSubscribe((subscription) -> {
-							Subscription prev = getSessionInfo(session).getSubscriptions().putIfAbsent(id, subscription);
-							if (prev != null) {
-								throw new SubscriptionExistsException();
-							}
+						SessionState state = this.sessionInfoMap.get(session.getId());
+						if (state == null) {
+							subscription.cancel();
+							return;
+						}
+						Subscription prev = state.getSubscriptions().putIfAbsent(id, subscription);
+						if (prev != null) {
+							throw new SubscriptionExistsException();
+						}
 					});
 		}
 		else {

spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java

@@ -38,6 +38,7 @@
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
 
 import org.springframework.aot.hint.RuntimeHints;
@@ -290,6 +291,25 @@ void messageAfterConnectionClosed() throws Exception {
 				this.handler.handleTextMessage(this.session, new TextMessage("{\"type\":\"ping\"}")));
 	}
 
+	@Test // gh-1501
+	void responseAfterConnectionClosed() throws Exception {
+		Sinks.Empty<Void> responseDelay = Sinks.empty();
+		GraphQlWebSocketHandler handler = initWebSocketHandler(
+				(request, chain) -> chain.next(request).delayUntil((response) -> responseDelay.asMono()));
+
+		handle(handler, new TextMessage("{\"type\":\"connection_init\"}"), new TextMessage(BOOK_SUBSCRIPTION));
+		handler.afterConnectionClosed(this.session, CloseStatus.NORMAL);
+
+		assertThatNoException().isThrownBy(responseDelay::tryEmitEmpty);
+
+		assertThat(this.session.isOpen()).isTrue();
+		StepVerifier.create(this.session.getOutput())
+				.consumeNextWith((message) -> assertMessageType(message, GraphQlWebSocketMessageType.CONNECTION_ACK))
+				.then(this.session::close) // Complete output Flux
+				.expectComplete()
+				.verify(TIMEOUT);
+	}
+
 	@Test
 	void connectionInitRejected() throws Exception {