Adding TraceId headers to responses in Spring REST webservice only works when response body is not empty

Question

We have a Spring Boot application, that contains this class, which should add as UUID as a header to each response to a requests to the app's REST operations.

/**
 * Adds the trace-id as a header to each response
 */
@ControllerAdvice
public class AddTraceIdHeaderToResponses implements ResponseBodyAdvice<Object> {
    @Override
    public boolean supports(@NonNull MethodParameter returnType, @NonNull Class converterType) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, @NonNull MethodParameter returnType, @NonNull     MediaType selectedContentType,
            @NonNull Class selectedConverterType, @NonNull ServerHttpRequest request, @NonNull     ServerHttpResponse response) {
        addTraceIdHeaderIfTraceIdIsNotNull(response);
        return body;
    }

    private static void addTraceIdHeaderIfTraceIdIsNotNull(ServerHttpResponse response) {
        String traceId = MDC.get("trace-id");
        if(traceId != null){
            response.getHeaders().add("trace-id", traceId);
        }
    }

}

This works well in the happy flow, when an operation returns a 200 response with a response body, but not if an exception occurs (for example, when a request is invalid and a 400 response is returned).

This app in question also has one operation that returns a 204 response (and no response body), and here the trace-id is missing as well.

How can I make sure the trace-id header is added to each response (so also if an exception occurs or the operation is not supposed to contain a response body)?

We tried implementing the ResponseBodyAdvice quoted above, and expected the trace-id header to be added to all responses. We found that this worked fine for 200 responses / our happy flow, but not when exceptions occur (or when a 204 is returned).