From db940e442547a6f8beda6ec47c3a533a9cf5f3a2 Mon Sep 17 00:00:00 2001
From: Zvonimir Sabljic <zvonimir@pythagora.io>
Date: Fri, 21 Feb 2025 10:15:59 +0100
Subject: [PATCH] Improved Sentry integration

---
 core/cli/main.py | 38 ++++++++++++++++++++++++++++++++++++++
 main.py          | 15 ---------------
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/core/cli/main.py b/core/cli/main.py
index 36528ecd..3789bc78 100644
--- a/core/cli/main.py
+++ b/core/cli/main.py
@@ -5,6 +5,14 @@ import sys
 from argparse import Namespace
 from asyncio import run
 
+try:
+    import sentry_sdk
+    from sentry_sdk.integrations.asyncio import AsyncioIntegration
+
+    SENTRY_AVAILABLE = True
+except ImportError:
+    SENTRY_AVAILABLE = False
+
 from core.agents.orchestrator import Orchestrator
 from core.cli.helpers import delete_project, init, list_projects, list_projects_json, load_project, show_config
 from core.config import LLMProvider, get_config
@@ -23,6 +31,21 @@ log = get_logger(__name__)
 telemetry_sent = False
 
 
+def init_sentry():
+    if SENTRY_AVAILABLE:
+        sentry_sdk.init(
+            dsn="https://4101633bc5560bae67d6eab013ba9686@o4508731634221056.ingest.us.sentry.io/4508732401909760",
+            send_default_pii=True,
+            traces_sample_rate=1.0,
+            integrations=[AsyncioIntegration()],
+        )
+
+
+def capture_exception(err):
+    if SENTRY_AVAILABLE:
+        sentry_sdk.capture_exception(err)
+
+
 async def cleanup(ui: UIBase):
     global telemetry_sent
     if not telemetry_sent:
@@ -62,6 +85,8 @@ async def run_project(sm: StateManager, ui: UIBase, args) -> bool:
         await sm.rollback()
     except APIError as err:
         log.warning(f"LLM API error occurred: {err.message}")
+        init_sentry()
+        capture_exception(err)
         await ui.send_message(
             f"Stopping Pythagora due to an error while calling the LLM API: {err.message}",
             source=pythagora_source,
@@ -70,6 +95,8 @@ async def run_project(sm: StateManager, ui: UIBase, args) -> bool:
         await sm.rollback()
     except CustomAssertionError as err:
         log.warning(f"Anthropic assertion error occurred: {str(err)}")
+        init_sentry()
+        capture_exception(err)
         await ui.send_message(
             f"Stopping Pythagora due to an error inside Anthropic SDK. {str(err)}",
             source=pythagora_source,
@@ -78,6 +105,9 @@ async def run_project(sm: StateManager, ui: UIBase, args) -> bool:
         await sm.rollback()
     except Exception as err:
         log.error(f"Uncaught exception: {err}", exc_info=True)
+        init_sentry()
+        capture_exception(err)
+
         stack_trace = telemetry.record_crash(err)
         await sm.rollback()
         await ui.send_message(
@@ -278,6 +308,9 @@ async def async_main(
 
     telemetry.set("user_contact", args.email)
 
+    if SENTRY_AVAILABLE and args.email:
+        sentry_sdk.set_user({"email": args.email})
+
     if args.extension_version:
         telemetry.set("is_extension", True)
         telemetry.set("extension_version", args.extension_version)
@@ -317,6 +350,11 @@ async def async_main(
 
     try:
         success = await run_pythagora_session(sm, ui, args)
+    except Exception as err:
+        log.error(f"Uncaught exception in main session: {err}", exc_info=True)
+        init_sentry()
+        capture_exception(err)
+        raise
     finally:
         await cleanup(ui)
 
diff --git a/main.py b/main.py
index 23dc3a54..c87287b7 100644
--- a/main.py
+++ b/main.py
@@ -2,21 +2,6 @@
 import os.path
 import sys
 
-try:
-    import sentry_sdk
-    from sentry_sdk.integrations.asyncio import AsyncioIntegration
-
-    sentry_sdk.init(
-        dsn="https://4101633bc5560bae67d6eab013ba9686@o4508731634221056.ingest.us.sentry.io/4508732401909760",
-        send_default_pii=True,
-        traces_sample_rate=1.0,
-        integrations=[AsyncioIntegration()],
-    )
-
-    sentry_sdk.profiler.start_profiler()
-except ImportError:
-    SENTRY_ENABLED = False
-
 try:
     from core.cli.main import run_pythagora
 except ImportError as err: