Quebra de linha em blocos de código para evitar barra de rolagem horizontal

É possível permitir a quebra de linha em blocos de código para evitar a barra de rolagem horizontal?

Ainda melhor se isso puder ser configurado para afetar apenas blocos de código de uma linguagem específica.

.lang-jai.hljs {
  /* Desabilitar max-height para evitar barra de rolagem vertical; sempre mostrar blocos de código completos */
  max-height: none;

  /* Habilitar quebra de linha no código jai para evitar barra de rolagem horizontal. */
  overflow-x: auto; 
  white-space: pre-wrap;
}