Add minimal SSE response support to azure-core

sdk/core/azure-core/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed synchronous streaming of non-replayable `BinaryData` response bodies.
+
 ### Other Changes
 
 ## 1.59.0 (2026-08-12)

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java

@@ -21,6 +21,7 @@
 import com.azure.core.implementation.util.ByteArrayContent;
 import com.azure.core.implementation.util.ByteBufferContent;
 import com.azure.core.implementation.util.HttpHeadersAccessHelper;
+import com.azure.core.implementation.util.HttpUtils;
 import com.azure.core.implementation.util.InputStreamContent;
 import com.azure.core.implementation.util.SerializableContent;
 import com.azure.core.implementation.util.StringContent;
@@ -506,8 +507,8 @@ private Long getAndLogContentLength(HttpHeaders headers, LoggingEventBuilder log
     /*
      * Determines if the request or response body should be logged.
      *
-     * <p>The request or response body is logged if the Content-Type is not "application/octet-stream" and the body
-     * isn't empty and is less than 16KB in size.</p>
+     * <p>The request or response body is logged if the Content-Type is neither "application/octet-stream" nor
+     * "text/event-stream" and the body isn't empty and is less than 16KB in size.</p>
      *
      * @param contentTypeHeader Content-Type header value.
      *
@@ -518,6 +519,7 @@ private Long getAndLogContentLength(HttpHeaders headers, LoggingEventBuilder log
     private static boolean shouldBodyBeLogged(String contentTypeHeader, Long contentLength) {
         return contentLength != null
             && !ContentType.APPLICATION_OCTET_STREAM.equalsIgnoreCase(contentTypeHeader)
+            && !HttpUtils.isTextEventStreamContentType(contentTypeHeader)
             && contentLength != 0
             && contentLength < MAX_BODY_LOG_SIZE;
     }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/FluxInputStream.java

@@ -29,7 +29,7 @@ public class FluxInputStream extends InputStream {
     private final Flux<ByteBuffer> data;
 
     // Subscription to request more data from as needed
-    private Subscription subscription;
+    private volatile Subscription subscription;
 
     private ByteArrayInputStream buffer;
 
@@ -141,6 +141,12 @@ public int read(byte[] b, int off, int len) throws IOException {
         }
     }
 
+    /**
+     * Closes the stream and cancels its Flux subscription. If the stream has not been read, closing subscribes and
+     * immediately cancels without requesting data so publisher cleanup can run.
+     *
+     * @throws IOException if the stream cannot be closed.
+     */
     @Override
     public void close() throws IOException {
         closed = true;
@@ -151,6 +157,9 @@ public void close() throws IOException {
         // Unblock any thread waiting in blockForData().
         lock.lock();
         try {
+            if (!subscribed) {
+                subscribeToData();
+            }
             waitingForData = false;
             dataAvailable.signal();
         } finally {
@@ -227,13 +236,13 @@ private void subscribeToData() {
                 this::signalOnCompleteOrError,
                 // Subscription consumer
                 subscription -> {
+                    this.subscription = subscription;
+                    this.subscribed = true;
                     if (this.closed) {
                         subscription.cancel();
                         return;
                     }
-                    this.subscription = subscription;
-                    this.subscribed = true;
-                    this.subscription.request(1);
+                    subscription.request(1);
                 });
     }
 

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/AsyncRestProxy.java

@@ -13,6 +13,7 @@
 import com.azure.core.http.rest.StreamResponse;
 import com.azure.core.implementation.TypeUtil;
 import com.azure.core.implementation.serializer.HttpResponseDecoder;
+import com.azure.core.implementation.util.HttpUtils;
 import com.azure.core.util.Base64Url;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
@@ -46,8 +47,6 @@
  */
 public class AsyncRestProxy extends RestProxyBase {
 
-    private static final String TEXT_EVENT_STREAM = "text/event-stream";
-
     /**
      * Create a RestProxy.
      *
@@ -143,23 +142,28 @@ private Mono<HttpResponseDecoder.HttpDecodedResponse> ensureExpectedStatus(
 
     private Mono<?> handleRestResponseReturnType(final HttpResponseDecoder.HttpDecodedResponse response,
         final SwaggerMethodParser methodParser, final Type entityType) {
+        final boolean isTextEventStream = HttpUtils.isTextEventStreamContentType(
+            response.getSourceResponse().getHeaders().getValue(HttpHeaderName.CONTENT_TYPE));
+        final ResponseBodyOwner responseBodyOwner
+            = isTextEventStream ? new ResponseBodyOwner(response.getSourceResponse()) : null;
         if (methodParser.isStreamResponse()) {
             return Mono.fromSupplier(() -> new StreamResponse(response.getSourceResponse()));
         } else if (TypeUtil.isTypeOrSubTypeOf(entityType, Response.class)) {
             final Type bodyType = TypeUtil.getRestResponseBodyType(entityType);
             if (TypeUtil.isTypeOrSubTypeOf(bodyType, Void.class)) {
-                return response.getSourceResponse()
-                    .getBody()
-                    .ignoreElements()
+                Flux<ByteBuffer> responseBody
+                    = responseBodyOwner == null ? response.getSourceResponse().getBody() : responseBodyOwner.getBody();
+                return responseBody.ignoreElements()
                     .then(Mono.fromCallable(() -> createResponse(response, entityType, null)));
             } else {
-                return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType)
-                    .map(bodyAsObject -> createResponse(response, entityType, bodyAsObject))
-                    .switchIfEmpty(Mono.fromCallable(() -> createResponse(response, entityType, null)));
+                return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, bodyType,
+                    responseBodyOwner).map(bodyAsObject -> createResponse(response, entityType, bodyAsObject))
+                        .switchIfEmpty(Mono.fromCallable(() -> createResponse(response, entityType, null)));
             }
         } else {
             // For now, we're just throwing if the Maybe didn't emit a value.
-            return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType);
+            return handleBodyReturnType(response.getSourceResponse(), decodeBytes(response), methodParser, entityType,
+                responseBodyOwner);
         }
     }
 
@@ -177,10 +181,17 @@ private static Function<byte[], Mono<Object>> decodeBytes(HttpResponseDecoder.Ht
     }
 
     static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[], Mono<Object>> getDecodedBody,
-        SwaggerMethodParser methodParser, Type entityType) {
+        SwaggerMethodParser methodParser, Type entityType, ResponseBodyOwner responseBodyOwner) {
         final int responseStatusCode = sourceResponse.getStatusCode();
         final HttpMethod httpMethod = methodParser.getHttpMethod();
         final Type returnValueWireType = methodParser.getReturnValueWireType();
+        if (responseBodyOwner == null
+            && HttpUtils
+                .isTextEventStreamContentType(sourceResponse.getHeaders().getValue(HttpHeaderName.CONTENT_TYPE))) {
+            responseBodyOwner = new ResponseBodyOwner(sourceResponse);
+        }
+        final Flux<ByteBuffer> responseBody
+            = responseBodyOwner == null ? sourceResponse.getBody() : responseBodyOwner.getBody();
 
         final Mono<?> asyncResult;
         if (httpMethod == HttpMethod.HEAD
@@ -199,20 +210,19 @@ static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[]
             asyncResult = responseBodyBytesAsync;
         } else if (FluxUtil.isFluxByteBuffer(entityType)) {
             // Mono<Flux<ByteBuffer>>
-            asyncResult = Mono.just(sourceResponse.getBody());
+            asyncResult = Mono.just(responseBody);
         } else if (TypeUtil.isTypeOrSubTypeOf(entityType, BinaryData.class)) {
-            String contentType = sourceResponse.getHeaders().getValue(HttpHeaderName.CONTENT_TYPE);
             // Mono<BinaryData>
             // The raw response is directly used to create an instance of BinaryData which then provides
             // different methods to read the response. The reading of the response is delayed until BinaryData
             // is read and depending on which format the content is converted into, the response is not necessarily
             // fully copied into memory resulting in lesser overall memory usage.
-            if (contentType != null && contentType.startsWith(TEXT_EVENT_STREAM)) {
-                // if the response content type is a stream, create a BinaryData instance with bufferContent set to
-                // false.
-                asyncResult = BinaryData.fromFlux(sourceResponse.getBody(), null, false);
+            if (responseBodyOwner != null) {
+                // If the response content type identifies a stream, create a BinaryData instance with bufferContent
+                // set to false.
+                asyncResult = BinaryData.fromFlux(responseBody, null, false);
             } else {
-                asyncResult = BinaryData.fromFlux(sourceResponse.getBody());
+                asyncResult = BinaryData.fromFlux(responseBody);
             }
         } else if (TypeUtil.isTypeOrSubTypeOf(entityType, InputStream.class)) {
             // Corresponds to the Open API 2.0 type "file" which is mapped to an InputStream.
@@ -224,6 +234,11 @@ static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[]
         return asyncResult;
     }
 
+    static Mono<?> handleBodyReturnType(HttpResponse sourceResponse, Function<byte[], Mono<Object>> getDecodedBody,
+        SwaggerMethodParser methodParser, Type entityType) {
+        return handleBodyReturnType(sourceResponse, getDecodedBody, methodParser, entityType, null);
+    }
+
     /**
      * Handle the provided asynchronous HTTP response and return the deserialized value.
      *
@@ -238,7 +253,6 @@ private Object handleRestReturnType(Mono<HttpResponseDecoder.HttpDecodedResponse
         EnumSet<ErrorOptions> errorOptionsSet) {
         final Mono<HttpResponseDecoder.HttpDecodedResponse> asyncExpectedResponse = endSpanWhenDone(
             ensureExpectedStatus(asyncHttpDecodedResponse, methodParser, options, errorOptionsSet), context);
-
         final Object result;
         if (TypeUtil.isTypeOrSubTypeOf(returnType, Mono.class)) {
             final Type monoTypeParam = TypeUtil.getTypeArgument(returnType);

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/RestProxyBase.java

@@ -35,13 +35,17 @@
 import com.azure.core.util.serializer.SerializerAdapter;
 import com.azure.core.util.tracing.Tracer;
 import reactor.core.Exceptions;
+import reactor.core.publisher.Flux;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Type;
 import java.net.URL;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.EnumSet;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 
 import static com.azure.core.util.FluxUtil.monoError;
@@ -213,6 +217,30 @@ public Response createResponse(HttpResponseDecoder.HttpDecodedResponse response,
         return RESPONSE_CONSTRUCTORS_CACHE.invoke(constructorReflectiveInvoker, response, bodyAsObject);
     }
 
+    static final class ResponseBodyOwner implements Closeable {
+        private final AtomicBoolean closed = new AtomicBoolean();
+        private final HttpResponse response;
+
+        ResponseBodyOwner(HttpResponse response) {
+            this.response = response;
+        }
+
+        Flux<ByteBuffer> getBody() {
+            return getBody(response.getBody());
+        }
+
+        Flux<ByteBuffer> getBody(Flux<ByteBuffer> responseBody) {
+            return Flux.using(() -> this, ignored -> responseBody, ResponseBodyOwner::close);
+        }
+
+        @Override
+        public void close() {
+            if (closed.compareAndSet(false, true)) {
+                response.close();
+            }
+        }
+    }
+
     /**
      * Starts the tracing span for the current service call, additionally set metadata attributes on the span by passing
      * additional context information.

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/SyncRestProxy.java

@@ -13,6 +13,8 @@
 import com.azure.core.implementation.ImplUtils;
 import com.azure.core.implementation.TypeUtil;
 import com.azure.core.implementation.serializer.HttpResponseDecoder;
+import com.azure.core.implementation.util.BinaryDataHelper;
+import com.azure.core.implementation.util.FluxByteBufferContent;
 import com.azure.core.util.Base64Url;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
@@ -145,7 +147,9 @@ private Object handleRestResponseReturnType(HttpResponseDecoder.HttpDecodedRespo
                 response.getSourceResponse().close();
                 return createResponse(response, entityType, null);
             } else {
-                Object bodyAsObject = handleBodyReturnType(response, methodParser, bodyType);
+                Object bodyAsObject = TypeUtil.isTypeOrSubTypeOf(bodyType, BinaryData.class)
+                    ? getOwnedResponseBody(response.getSourceResponse())
+                    : handleBodyReturnType(response, methodParser, bodyType);
                 Response<?> httpResponse = createResponse(response, entityType, bodyAsObject);
                 if (httpResponse == null) {
                     return createResponse(response, entityType, null);
@@ -195,6 +199,17 @@ private Object handleBodyReturnType(HttpResponseDecoder.HttpDecodedResponse resp
         return result;
     }
 
+    private static BinaryData getOwnedResponseBody(HttpResponse response) {
+        BinaryData responseBody = response.getBodyAsBinaryData();
+        if (responseBody == null || responseBody.isReplayable()) {
+            return responseBody;
+        }
+
+        ResponseBodyOwner responseBodyOwner = new ResponseBodyOwner(response);
+        return BinaryDataHelper.createBinaryData(new FluxByteBufferContent(
+            responseBodyOwner.getBody(responseBody.toFluxByteBuffer()), responseBody.getLength(), false));
+    }
+
     /**
      * Handle the provided asynchronous HTTP response and return the deserialized value.
      *

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxByteBufferContent.java

@@ -3,6 +3,7 @@
 
 package com.azure.core.implementation.util;
 
+import com.azure.core.implementation.FluxInputStream;
 import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.serializer.ObjectSerializer;
@@ -100,9 +101,24 @@ public <T> T toObject(TypeReference<T> typeReference, ObjectSerializer serialize
         return serializer.deserializeFromBytes(toBytes(), typeReference);
     }
 
+    /**
+     * Returns an in-memory stream for cached or replayable content. Uncached non-replayable content is streamed
+     * incrementally without buffering the full Flux.
+     *
+     * @return A stream over this content.
+     */
     @Override
     public InputStream toStream() {
-        return new ByteArrayInputStream(toBytes());
+        byte[] cachedBytes = BYTES_UPDATER.get(this);
+        if (cachedBytes != null) {
+            return new ByteArrayInputStream(cachedBytes);
+        }
+
+        if (isReplayable) {
+            return new ByteArrayInputStream(toBytes());
+        }
+
+        return new FluxInputStream(content);
     }
 
     @Override

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/HttpUtils.java

@@ -6,6 +6,9 @@
 import com.azure.core.util.logging.ClientLogger;
 
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 
 import static com.azure.core.util.Configuration.PROPERTY_AZURE_REQUEST_CONNECT_TIMEOUT;
 import static com.azure.core.util.Configuration.PROPERTY_AZURE_REQUEST_READ_TIMEOUT;
@@ -17,6 +20,7 @@
  * Utilities shared with HttpClient implementations.
  */
 public final class HttpUtils {
+    private static final String TEXT_EVENT_STREAM = "text/event-stream";
     private static final ClientLogger LOGGER = new ClientLogger(HttpUtils.class);
 
     private static final Duration MINIMUM_TIMEOUT = Duration.ofMillis(1);
@@ -60,6 +64,56 @@ public final class HttpUtils {
      */
     public static final String AZURE_EAGERLY_CONVERT_HEADERS = "azure-eagerly-convert-headers";
 
+    /**
+     * Determines whether a Content-Type header identifies exactly one {@code text/event-stream} representation.
+     * Charset parameters don't affect this determination as event streams are always decoded as UTF-8.
+     *
+     * @param headerValue The Content-Type header value.
+     * @return Whether the header identifies a {@code text/event-stream} representation.
+     */
+    public static boolean isTextEventStreamContentType(String headerValue) {
+        if (headerValue == null) {
+            return false;
+        }
+
+        List<String> mediaTypeAndParameters = splitHeaderValue(headerValue, ';');
+        if (mediaTypeAndParameters.size() == 0
+            || !TEXT_EVENT_STREAM.equalsIgnoreCase(mediaTypeAndParameters.get(0).trim())
+            || splitHeaderValue(headerValue, ',').size() != 1) {
+            return false;
+        }
+
+        return true;
+    }
+
+    private static List<String> splitHeaderValue(String value, char delimiter) {
+        List<String> segments = new ArrayList<>();
+        int start = 0;
+        boolean quoted = false;
+        boolean escaped = false;
+
+        for (int i = 0; i < value.length(); i++) {
+            char character = value.charAt(i);
+            if (escaped) {
+                escaped = false;
+            } else if (quoted && character == '\\') {
+                escaped = true;
+            } else if (character == '"') {
+                quoted = !quoted;
+            } else if (!quoted && character == delimiter) {
+                segments.add(value.substring(start, i));
+                start = i + 1;
+            }
+        }
+
+        if (quoted) {
+            return Collections.emptyList();
+        }
+
+        segments.add(value.substring(start));
+        return segments;
+    }
+
     /**
      * Gets the default connect timeout.
      *