def test_simple_chat_engine_astream_exception_handling():
    """
    Test that an exception thrown while retrieving the streamed LLM response gets bubbled up to the user.
    Also tests that the non-retrieved exception does not remain in an task that was not awaited leading to
    a 'Task exception was never retrieved' message during garbage collection.
    """

    class ExceptionThrownInTest(Exception):
        pass

    class ExceptionMockLLM(MockLLM):
        """Raises an exception while streaming back the mocked LLM response."""

        @classmethod
        def class_name(cls) -> str:
            return "ExceptionMockLLM"

        @llm_completion_callback()
        def stream_complete(
            self, prompt: str, formatted: bool = False, **kwargs: Any
        ) -> CompletionResponseGen:
            def gen_prompt() -> CompletionResponseGen:
                for ch in prompt:
                    yield CompletionResponse(
                        text=prompt,
                        delta=ch,
                    )
                raise ExceptionThrownInTest("Exception thrown for testing purposes")

            return gen_prompt()

    async def async_part():
        engine = SimpleChatEngine.from_defaults(
            llm=ExceptionMockLLM(), memory=ChatMemoryBuffer.from_defaults()
        )
        response = await engine.astream_chat("Hello World!")

        with pytest.raises(ExceptionThrownInTest):
            async for response_part in response.async_response_gen():
                pass

    not_retrieved_exception = False

    def custom_exception_handler(loop, context):
        if context.get("message") == "Task exception was never retrieved":
            nonlocal not_retrieved_exception
            not_retrieved_exception = True

    loop = asyncio.new_event_loop()
    loop.set_exception_handler(custom_exception_handler)
    result = loop.run_until_complete(async_part())
    loop.close()
    gc.collect()
    if not_retrieved_exception:
        pytest.fail(
            "Exception was not correctly handled - ended up in asyncio cleanup performed during garbage collection"
        )

    Test that an exception thrown while retrieving the streamed LLM response gets bubbled up to the user.
    Also tests that the non-retrieved exception does not remain in an task that was not awaited leading to
    a 'Task exception was never retrieved' message during garbage collection.
    
llama-index-core/tests/chat_engine/test_simple.py