Update WebFlux fragment stream rendering

Thymeleaf has its own special handling for SSE that gets in the way
of fragment rendering. This is why we need to set the response
content-type before streaming, and then pass text/html to the
View for rendering each fragment.

See gh-33194

Author: rstoyanchev

spring-webflux/src/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java

@@ -267,20 +267,24 @@ else if (Rendering.class.isAssignableFrom(clazz)) {
 								Mono.just(Collections.singletonList((View) view)));
 					}
 					else if (FragmentsRendering.class.isAssignableFrom(clazz)) {
+						ServerHttpResponse response = exchange.getResponse();
 						FragmentsRendering render = (FragmentsRendering) returnValue;
 						HttpStatusCode status = render.status();
 						if (status != null) {
-							exchange.getResponse().setStatusCode(status);
+							response.setStatusCode(status);
 						}
-						exchange.getResponse().getHeaders().putAll(render.headers());
+						response.getHeaders().putAll(render.headers());
 						bindingContext.updateModel(exchange);
 
 						StreamHandler streamHandler = getStreamHandler(exchange);
+						if (streamHandler != null) {
+							streamHandler.updateResponse(exchange);
+						}
 
 						Flux<Flux<DataBuffer>> renderFlux = render.fragments().concatMap(fragment ->
 								renderFragment(fragment, streamHandler, locale, bindingContext, exchange));
 
-						return exchange.getResponse().writeAndFlushWith(renderFlux);
+						return response.writeAndFlushWith(renderFlux);
 					}
 					else if (Model.class.isAssignableFrom(clazz)) {
 						model.addAllAttributes(((Model) returnValue).asMap());
@@ -299,7 +303,7 @@ else if (View.class.isAssignableFrom(clazz)) {
 						viewsMono = resolveViews(getDefaultViewName(exchange), locale);
 					}
 					bindingContext.updateModel(exchange);
-					return viewsMono.flatMap(views -> render(views, model.asMap(), bindingContext, exchange));
+					return viewsMono.flatMap(views -> render(views, model.asMap(), null, bindingContext, exchange));
 				});
 	}
 
@@ -346,10 +350,16 @@ private Mono<Flux<DataBuffer>> renderFragment(
 				Mono.just(List.of(fragment.view())) :
 				resolveViews(fragment.viewName() != null ? fragment.viewName() : getDefaultViewName(exchange), locale));
 
-		return selectedViews.flatMap(views -> render(views, fragment.model(), bindingContext, mutatedExchange))
-				.then(Mono.fromSupplier(() -> (streamHandler != null ?
-						streamHandler.format(response.getBodyFlux(), fragment, exchange) :
-						response.getBodyFlux())));
+		Map<String, Object> model = fragment.model();
+
+		if (streamHandler != null) {
+			return selectedViews.flatMap(views -> render(views, model, MediaType.TEXT_HTML, bindingContext, mutatedExchange))
+					.then(Mono.fromSupplier(() -> streamHandler.format(response.getBodyFlux(), fragment, exchange)));
+		}
+		else {
+			return selectedViews.flatMap(views -> render(views, model, null, bindingContext, mutatedExchange))
+					.then(Mono.fromSupplier(response::getBodyFlux));
+		}
 	}
 
 	@Nullable
@@ -369,7 +379,8 @@ private String getNameForReturnValue(MethodParameter returnType) {
 				.orElseGet(() -> Conventions.getVariableNameForParameter(returnType));
 	}
 
-	private Mono<? extends Void> render(List<View> views, Map<String, Object> model,
+	private Mono<? extends Void> render(
+			List<View> views, Map<String, Object> model, @Nullable MediaType bestMediaType,
 			BindingContext bindingContext, ServerWebExchange exchange) {
 
 		for (View view : views) {
@@ -378,19 +389,20 @@ private Mono<? extends Void> render(List<View> views, Map<String, Object> model,
 			}
 		}
 		List<MediaType> mediaTypes = getMediaTypes(views);
-		MediaType bestMediaType;
-		try {
-			bestMediaType = selectMediaType(exchange, () -> mediaTypes);
-		}
-		catch (NotAcceptableStatusException ex) {
-			HttpStatusCode statusCode = exchange.getResponse().getStatusCode();
-			if (statusCode != null && statusCode.isError()) {
-				if (logger.isDebugEnabled()) {
-					logger.debug("Ignoring error response content (if any). " + ex.getReason());
+		if (bestMediaType == null) {
+			try {
+				bestMediaType = selectMediaType(exchange, () -> mediaTypes);
+			}
+			catch (NotAcceptableStatusException ex) {
+				HttpStatusCode statusCode = exchange.getResponse().getStatusCode();
+				if (statusCode != null && statusCode.isError()) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("Ignoring error response content (if any). " + ex.getReason());
+					}
+					return Mono.empty();
 				}
-				return Mono.empty();
+				throw ex;
 			}
-			throw ex;
 		}
 		if (bestMediaType != null) {
 			for (View view : views) {
@@ -427,15 +439,23 @@ private static class BodySavingResponse extends ServerHttpResponseDecorator {
 		@Nullable
 		private Flux<DataBuffer> bodyFlux;
 
-		private final HttpHeaders headers;
+		@Nullable
+		private HttpHeaders headers;
 
 		BodySavingResponse(ServerHttpResponse delegate) {
 			super(delegate);
-			this.headers = new HttpHeaders(delegate.getHeaders()); // Ignore header changes
 		}
 
 		@Override
 		public HttpHeaders getHeaders() {
+			if (!super.getHeaders().containsKey(HttpHeaders.CONTENT_TYPE)) {
+				return super.getHeaders();
+			}
+			// Content-type is set, ignore further updates
+			if (this.headers == null) {
+				this.headers = new HttpHeaders();
+				this.headers.putAll(super.getHeaders());
+			}
 			return this.headers;
 		}
 
@@ -468,6 +488,11 @@ private interface StreamHandler {
 		 */
 		boolean supports(ServerHttpRequest request);
 
+		/**
+		 * Update the response before streaming, e.g. to set the content-type.
+		 */
+		void updateResponse(ServerWebExchange exchange);
+
 		/**
 		 * Format the given fragment.
 		 * @param fragmentContent the fragment serialized to data buffers
@@ -476,7 +501,6 @@ private interface StreamHandler {
 		 * @return the formatted fragment
 		 */
 		Flux<DataBuffer> format(Flux<DataBuffer> fragmentContent, Fragment fragment, ServerWebExchange exchange);
-
 	}
 
 
@@ -492,20 +516,14 @@ public boolean supports(ServerHttpRequest request) {
 		}
 
 		@Override
-		public Flux<DataBuffer> format(
-				Flux<DataBuffer> fragmentContent, Fragment fragment, ServerWebExchange exchange) {
-
+		public void updateResponse(ServerWebExchange exchange) {
+			MediaType mediaType = MediaType.TEXT_EVENT_STREAM;
 			Charset charset = getCharset(exchange.getRequest());
-			DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
-
-			String eventLine = fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : "";
-
-			return Flux.concat(
-					Flux.just(encodeText(eventLine + "data:", charset, bufferFactory)),
-					fragmentContent,
-					Flux.just(encodeText("\n\n", charset, bufferFactory)));
+			mediaType = (charset != null ? new MediaType(mediaType, charset) : mediaType);
+			exchange.getResponse().getHeaders().setContentType(mediaType);
 		}
 
+		@Nullable
 		private Charset getCharset(ServerHttpRequest request) {
 			for (MediaType mediaType : request.getHeaders().getAccept()) {
 				if (mediaType.isCompatibleWith(MediaType.TEXT_EVENT_STREAM)) {
@@ -515,7 +533,26 @@ private Charset getCharset(ServerHttpRequest request) {
 					break;
 				}
 			}
-			return StandardCharsets.UTF_8;
+			return null;
+		}
+
+		@Override
+		public Flux<DataBuffer> format(
+				Flux<DataBuffer> fragmentContent, Fragment fragment, ServerWebExchange exchange) {
+
+			Charset charset = StandardCharsets.UTF_8;
+			MediaType contentType = exchange.getResponse().getHeaders().getContentType();
+			if (contentType != null && contentType.getCharset() != null) {
+				charset = contentType.getCharset();
+			}
+
+			DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
+
+			String eventLine = fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : "";
+			DataBuffer prefix = encodeText(eventLine + "data:", charset, bufferFactory);
+			DataBuffer suffix = encodeText("\n\n", charset, bufferFactory);
+
+			return Flux.concat(Flux.just(prefix), fragmentContent, Flux.just(suffix));
 		}
 
 		private DataBuffer encodeText(String text, Charset charset, DataBufferFactory bufferFactory) {