From af1f79afb7937a0009e553a141b031e0329c4231 Mon Sep 17 00:00:00 2001 From: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:08:14 +0000 Subject: [PATCH] FIX: Move streaming PCM file work off the event loop in PromptNormalizer --- pyrit/prompt_normalizer/prompt_normalizer.py | 56 ++++++++- .../test_prompt_normalizer.py | 118 ++++++++++++++++++ 2 files changed, 171 insertions(+), 3 deletions(-) diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 2ff851811f..c1456513ca 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -392,7 +392,7 @@ async def convert_audio_async( if not converter_configurations: return raw_pcm - input_path = _write_pcm_to_temp_wav( + input_path = await _write_pcm_to_temp_wav_async( raw_pcm=raw_pcm, sample_rate_hz=sample_rate_hz, num_channels=num_channels, @@ -411,7 +411,9 @@ async def convert_audio_async( converter_configurations=converter_configurations, message=message, ) - actual_rate, actual_channels, actual_width, converted_pcm = _read_pcm_from_wav(piece.converted_value) + actual_rate, actual_channels, actual_width, converted_pcm = await asyncio.to_thread( + _read_pcm_from_wav, piece.converted_value + ) if (actual_rate, actual_channels, actual_width) != ( sample_rate_hz, num_channels, @@ -425,7 +427,8 @@ async def convert_audio_async( ) return converted_pcm finally: - Path(input_path).unlink(missing_ok=True) + # The worker thread finishes the removal even if this await is cancelled. + await asyncio.to_thread(Path(input_path).unlink, missing_ok=True) async def _calc_hash_async(self, request: Message) -> None: """Add a request to the memory.""" @@ -500,6 +503,53 @@ async def add_prepended_conversation_to_memory_async( return prepended_conversation +async def _write_pcm_to_temp_wav_async( + *, + raw_pcm: bytes, + sample_rate_hz: int, + num_channels: int, + sample_width_bytes: int, +) -> str: + """ + Write the PCM to a new temp WAV file in a worker thread and return its path. + + A worker thread cannot be stopped, so a cancelled caller would never learn the path of the + file the thread goes on to create. In that case the file is removed once the write finishes. + + Args: + raw_pcm (bytes): Raw PCM audio samples (no WAV header). + sample_rate_hz (int): Sample rate of the PCM in Hz. + num_channels (int): Channel count. + sample_width_bytes (int): Bytes per sample. + + Returns: + str: The path of the new temp WAV file. + + Raises: + asyncio.CancelledError: If the caller is cancelled before the write finishes. + """ + write = asyncio.ensure_future( + asyncio.to_thread( + _write_pcm_to_temp_wav, + raw_pcm=raw_pcm, + sample_rate_hz=sample_rate_hz, + num_channels=num_channels, + sample_width_bytes=sample_width_bytes, + ) + ) + try: + return await asyncio.shield(write) + except asyncio.CancelledError: + write.add_done_callback(_remove_abandoned_temp_wav) + raise + + +def _remove_abandoned_temp_wav(write: asyncio.Future[str]) -> None: + """Remove the temp WAV a cancelled caller left behind, if the write produced one.""" + if not write.cancelled() and write.exception() is None: + Path(write.result()).unlink(missing_ok=True) + + def _write_pcm_to_temp_wav( *, raw_pcm: bytes, diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index e06558b44c..9581f79a27 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -5,6 +5,7 @@ import codecs import os import tempfile +import threading import wave from collections.abc import AsyncGenerator from contextlib import asynccontextmanager @@ -45,6 +46,7 @@ SeedPrompt, ) from pyrit.prompt_normalizer import NormalizerRequest, PromptNormalizer +from pyrit.prompt_normalizer import prompt_normalizer as prompt_normalizer_module from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) @@ -1173,3 +1175,119 @@ async def swap_to_stereo(*, converter_configurations, message): ) finally: Path(wrong_channels_path).unlink(missing_ok=True) + + +async def test_convert_audio_async_runs_file_work_off_the_event_loop( + mock_memory_instance, sample_pcm, dummy_audio_converter_config +): + event_loop_thread = threading.get_ident() + file_work_threads: dict[str, int] = {} + + def record_thread(name, function): + def run(*args, **kwargs): + file_work_threads[name] = threading.get_ident() + return function(*args, **kwargs) + + return run + + normalizer = PromptNormalizer() + with ( + patch.object( + prompt_normalizer_module, + "_write_pcm_to_temp_wav", + side_effect=record_thread("write", prompt_normalizer_module._write_pcm_to_temp_wav), + ), + patch.object( + prompt_normalizer_module, + "_read_pcm_from_wav", + side_effect=record_thread("read", prompt_normalizer_module._read_pcm_from_wav), + ), + patch.object(Path, "unlink", autospec=True, side_effect=record_thread("unlink", Path.unlink)), + patch.object(normalizer, "convert_values_async", new_callable=AsyncMock), + ): + result = await normalizer.convert_audio_async( + raw_pcm=sample_pcm, + converter_configurations=[dummy_audio_converter_config], + sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ, + num_channels=_AUDIO_NUM_CHANNELS, + sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES, + ) + + assert result == sample_pcm + assert set(file_work_threads) == {"write", "read", "unlink"} + assert event_loop_thread not in file_work_threads.values() + + +async def test_convert_audio_async_cleans_up_temp_file_when_cancelled_during_conversion( + mock_memory_instance, sample_pcm, dummy_audio_converter_config +): + captured_paths: list[str] = [] + converting = asyncio.Event() + + async def capture_then_wait(*, converter_configurations, message): + captured_paths.append(message.message_pieces[0].converted_value) + converting.set() + await asyncio.Event().wait() + + normalizer = PromptNormalizer() + with patch.object(normalizer, "convert_values_async", side_effect=capture_then_wait): + task = asyncio.create_task( + normalizer.convert_audio_async( + raw_pcm=sample_pcm, + converter_configurations=[dummy_audio_converter_config], + sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ, + num_channels=_AUDIO_NUM_CHANNELS, + sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES, + ) + ) + await converting.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert len(captured_paths) == 1 + assert not Path(captured_paths[0]).exists() + + +async def test_convert_audio_async_removes_temp_file_written_after_cancellation( + mock_memory_instance, sample_pcm, dummy_audio_converter_config +): + write_started = threading.Event() + release_write = threading.Event() + written_paths: list[str] = [] + write_pcm_to_temp_wav = prompt_normalizer_module._write_pcm_to_temp_wav + + def write_after_release(**kwargs): + write_started.set() + release_write.wait(timeout=5) + written_paths.append(write_pcm_to_temp_wav(**kwargs)) + return written_paths[-1] + + normalizer = PromptNormalizer() + with ( + patch.object(prompt_normalizer_module, "_write_pcm_to_temp_wav", side_effect=write_after_release), + patch.object(normalizer, "convert_values_async", new_callable=AsyncMock) as mock_convert, + ): + task = asyncio.create_task( + normalizer.convert_audio_async( + raw_pcm=sample_pcm, + converter_configurations=[dummy_audio_converter_config], + sample_rate_hz=_AUDIO_SAMPLE_RATE_HZ, + num_channels=_AUDIO_NUM_CHANNELS, + sample_width_bytes=_AUDIO_SAMPLE_WIDTH_BYTES, + ) + ) + assert await asyncio.to_thread(write_started.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + release_write.set() + for _ in range(500): + if written_paths and not Path(written_paths[0]).exists(): + break + await asyncio.sleep(0.01) + + assert len(written_paths) == 1 + assert not Path(written_paths[0]).exists() + mock_convert.assert_not_called()