Ajuste de línea en bloques de código para evitar la barra de desplazamiento horizontal

¿Es posible permitir el ajuste de línea en los bloques de código para evitar la barra de desplazamiento horizontal?

Aún mejor si esto se puede configurar para afectar solo a los bloques de código de un lenguaje específico.

.lang-jai.hljs {
  /* Desactivar max-height para evitar la barra de desplazamiento vertical; siempre mostrar los bloques de código completos */
  max-height: none;

  /* Habilitar el ajuste de línea en el código jai para evitar la barra de desplazamiento horizontal. */
  overflow-x: auto; 
  white-space: pre-wrap;
}