Display exception (AbanteAI#250)

Co-authored-by: Wayde Gilliam <[email protected]>
Co-authored-by: BioBootloader <[email protected]>

mentat/session.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import traceback
 from asyncio import Task
 from pathlib import Path
 from typing import List, Optional
@@ -119,6 +120,8 @@ async def run_main():
                 await self.stop()
             except asyncio.CancelledError:
                 pass
+            except Exception:
+                traceback.print_exc()
 
         setup_logging()
         self._main_task: Task[None] = asyncio.create_task(run_main())

mentat/terminal/client.py

@@ -111,6 +111,10 @@ def _init_signal_handlers(self):
         signal.signal(signal.SIGINT, self._handle_sig_int)
 
     async def _startup(self):
+        def session_start_callback(_: asyncio.Task[None]):
+            """Shutdown the Terminal Client if the Session stops"""
+            self._should_exit.set()
+
         self.session = Session(
             self.paths,
             self.exclude_paths,
@@ -119,7 +123,8 @@ async def _startup(self):
             self.pr_diff,
             self.config,
         )
-        self.session.start()
+        session_start_task = self.session.start()
+        session_start_task.add_done_callback(session_start_callback)
         # Logging is setup in session.start()
         logging.debug("Running startup")
 

tests/system_test.py

@@ -53,6 +53,25 @@ async def test_system(mock_call_llm_api, mock_setup_api_key, mock_collect_user_i
     assert content == expected_content
 
 
+@pytest.mark.asyncio
+async def test_system_exits_on_exception(
+    mock_call_llm_api, mock_setup_api_key, mock_collect_user_input
+):
+    mock_collect_user_input.set_stream_messages(
+        [
+            "respond with 'hello'",
+        ]
+    )
+
+    # if we don't catch this and shutdown properly, pytest will fail test
+    # with "Task was destroyed but it is pending!"
+    mock_call_llm_api.side_effect = Exception("Something went wrong")
+
+    session = Session()
+    await session.start()
+    session.stream.stop()
+
+
 @pytest.mark.asyncio
 async def test_interactive_change_selection(
     mock_call_llm_api, mock_setup_api_key, mock_collect_user_input