Log more info about internal errors (#346)

Author: alexmojaki

logfire-api/logfire_api/_internal/stack_info.pyi

@@ -7,7 +7,7 @@ STACK_INFO_KEYS: Incomplete
 SITE_PACKAGES_DIR: Incomplete
 PYTHON_LIB_DIR: Incomplete
 LOGFIRE_DIR: Incomplete
-PREFIXES: Incomplete
+NON_USER_CODE_PREFIXES: Incomplete
 
 def get_filepath_attribute(file: str) -> StackInfo: ...
 def get_code_object_info(code: CodeType) -> StackInfo: ...

logfire-api/logfire_api/_internal/utils.pyi

@@ -1,5 +1,6 @@
 from _typeshed import Incomplete
 from collections.abc import Generator
+from logfire._internal.stack_info import is_user_code as is_user_code
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import Event as Event, ReadableSpan

logfire/_internal/main.py

@@ -8,7 +8,6 @@
 import warnings
 from functools import cached_property, partial
 from time import time
-from types import TracebackType
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Literal, Sequence, TypeVar, Union, cast
 
 import opentelemetry.context as context_api
@@ -52,7 +51,7 @@
 from .metrics import ProxyMeterProvider
 from .stack_info import get_user_stack_info
 from .tracer import ProxyTracerProvider
-from .utils import handle_internal_errors, log_internal_error, uniquify_sequence
+from .utils import SysExcInfo, handle_internal_errors, log_internal_error, uniquify_sequence
 
 if TYPE_CHECKING:
     import anthropic
@@ -88,13 +87,7 @@
 # 1. It's convenient to pass the result of sys.exc_info() directly
 # 2. It mirrors the exc_info argument of the stdlib logging methods
 # 3. The argument name exc_info is very suggestive of the sys function.
-ExcInfo: typing.TypeAlias = Union[
-    'tuple[type[BaseException], BaseException, TracebackType | None]',
-    'tuple[None, None, None]',
-    BaseException,
-    bool,
-    None,
-]
+ExcInfo: typing.TypeAlias = Union[SysExcInfo, BaseException, bool, None]
 
 
 class Logfire:

logfire/_internal/stack_info.py

@@ -22,7 +22,7 @@
 SITE_PACKAGES_DIR = str(Path(opentelemetry.sdk.trace.__file__).parent.parent.parent.parent.absolute())
 PYTHON_LIB_DIR = str(Path(inspect.__file__).parent.absolute())
 LOGFIRE_DIR = str(Path(logfire.__file__).parent.absolute())
-PREFIXES = (SITE_PACKAGES_DIR, PYTHON_LIB_DIR, LOGFIRE_DIR)
+NON_USER_CODE_PREFIXES = (SITE_PACKAGES_DIR, PYTHON_LIB_DIR, LOGFIRE_DIR)
 
 
 def get_filepath_attribute(file: str) -> StackInfo:
@@ -95,7 +95,7 @@ def is_user_code(code: CodeType) -> bool:
         On the other hand, generator expressions and lambdas might be called far away from where they are defined.
     """
     return not (
-        str(Path(code.co_filename).absolute()).startswith(PREFIXES)
+        str(Path(code.co_filename).absolute()).startswith(NON_USER_CODE_PREFIXES)
         or code.co_name in ('<listcomp>', '<dictcomp>', '<setcomp>')
     )
 

logfire/_internal/utils.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import inspect
 import json
 import logging
 import os
 import sys
+import typing
 from contextlib import contextmanager
 from pathlib import Path
+from types import TracebackType
 from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Sequence, Tuple, TypedDict, TypeVar, Union
 
 from opentelemetry import context, trace as trace_api
@@ -16,6 +19,8 @@
 from opentelemetry.util import types as otel_types
 from requests import RequestException, Response
 
+from logfire._internal.stack_info import is_user_code
+
 if TYPE_CHECKING:
     from packaging.version import Version
 
@@ -252,7 +257,81 @@ def log_internal_error():
         raise
 
     with suppress_instrumentation():  # prevent infinite recursion from the logging integration
-        logger.exception('Internal error in Logfire')
+        logger.exception('Internal error in Logfire', exc_info=_internal_error_exc_info())
+
+
+SysExcInfo: typing.TypeAlias = Union[
+    'tuple[type[BaseException], BaseException, TracebackType | None]',
+    'tuple[None, None, None]',
+]
+"""
+The return type of sys.exc_info(): exc_type, exc_val, exc_tb.
+"""
+
+
+def _internal_error_exc_info() -> SysExcInfo:
+    """Returns an exc_info tuple with a nicely tweaked traceback."""
+    original_exc_info: SysExcInfo = sys.exc_info()
+    exc_type, exc_val, original_tb = original_exc_info
+    try:
+        # First remove redundant frames already in the traceback about where the error was raised.
+        tb = original_tb
+        if tb and tb.tb_frame and tb.tb_frame.f_code is _HANDLE_INTERNAL_ERRORS_CODE:
+            # Skip the 'yield' line in _handle_internal_errors
+            tb = tb.tb_next
+
+        if (
+            tb
+            and tb.tb_frame
+            and tb.tb_frame.f_code.co_filename == contextmanager.__code__.co_filename
+            and tb.tb_frame.f_code.co_name == 'inner'
+        ):
+            # Skip the 'inner' function frame when handle_internal_errors is used as a decorator.
+            # It looks like `return func(*args, **kwds)`
+            tb = tb.tb_next
+
+        # Now add useful outer frames that give context, but skipping frames that are just about handling the error.
+        frame = inspect.currentframe()
+        # Skip this frame right here.
+        assert frame
+        frame = frame.f_back
+
+        if frame and frame.f_code is log_internal_error.__code__:  # pragma: no branch
+            # This function is always called from log_internal_error, so skip that frame.
+            frame = frame.f_back
+            assert frame
+
+            if frame.f_code is _HANDLE_INTERNAL_ERRORS_CODE:
+                # Skip the line in _handle_internal_errors that calls log_internal_error
+                frame = frame.f_back
+                # Skip the frame defining the _handle_internal_errors context manager
+                assert frame and frame.f_code.co_name == '__exit__'
+                frame = frame.f_back
+                assert frame
+                # Skip the frame calling the context manager, on the `with` line.
+                frame = frame.f_back
+            else:
+                # `log_internal_error()` was called directly, so just skip that frame. No context manager stuff.
+                frame = frame.f_back
+
+        # Now add all remaining frames from internal logfire code.
+        while frame and not is_user_code(frame.f_code):
+            tb = TracebackType(tb_next=tb, tb_frame=frame, tb_lasti=frame.f_lasti, tb_lineno=frame.f_lineno)
+            frame = frame.f_back
+
+        # Add up to 3 frames from user code.
+        for _ in range(3):
+            if not frame:  # pragma: no cover
+                break
+            tb = TracebackType(tb_next=tb, tb_frame=frame, tb_lasti=frame.f_lasti, tb_lineno=frame.f_lineno)
+            frame = frame.f_back
+
+        assert exc_type
+        assert exc_val
+        exc_val = exc_val.with_traceback(tb)
+        return exc_type, exc_val, tb
+    except Exception:  # pragma: no cover
+        return original_exc_info
 
 
 @contextmanager
@@ -263,6 +342,9 @@ def handle_internal_errors():
         log_internal_error()
 
 
+_HANDLE_INTERNAL_ERRORS_CODE = inspect.unwrap(handle_internal_errors).__code__
+
+
 def maybe_capture_server_headers(capture: bool):
     if capture:
         os.environ['OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST'] = '.*'

tests/import_used_for_tests/internal_error_handling/__init__.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+from logfire._internal.utils import handle_internal_errors, log_internal_error
+
+
+def inner1():
+    raise ValueError('inner1')
+
+
+def inner2():
+    inner1()
+
+
+@handle_internal_errors()
+def using_decorator():
+    inner2()
+
+
+def using_context_manager():
+    with handle_internal_errors():
+        inner2()
+
+
+def using_try_except():
+    try:
+        inner2()
+    except Exception:
+        log_internal_error()
+
+
+def outer1(func: Any):
+    func()
+
+
+def outer2(func: Any):
+    outer1(func)

tests/import_used_for_tests/internal_error_handling/internal_logfire_code_example.py

@@ -0,0 +1,32 @@
+from tests.import_used_for_tests.internal_error_handling.internal_logfire_code_example import (
+    outer2,
+    using_context_manager,
+    using_decorator,
+    using_try_except,
+)
+
+
+def user1():
+    user2()
+
+
+def user2():
+    user3()
+
+
+def user3():
+    user4()
+
+
+def user4():
+    user5()
+
+
+def user5():
+    user6()
+
+
+def user6():
+    outer2(using_decorator)
+    outer2(using_context_manager)
+    outer2(using_try_except)

tests/import_used_for_tests/internal_error_handling/user_code_example.py

@@ -1,8 +1,11 @@
 import pytest
 import requests
 import requests_mock
+from inline_snapshot import snapshot
 
+import logfire._internal.stack_info
 from logfire._internal.utils import UnexpectedResponse, handle_internal_errors
+from tests.import_used_for_tests.internal_error_handling import internal_logfire_code_example, user_code_example
 
 
 def test_raise_for_status() -> None:
@@ -21,3 +24,94 @@ def test_reraise_internal_exception():
     with pytest.raises(ZeroDivisionError):
         with handle_internal_errors():
             str(1 / 0)
+
+
+def test_internal_exception_tb(caplog: pytest.LogCaptureFixture):
+    # Pretend that `internal_logfire_code_example` is a module within logfire,
+    # so all frames from it should be included.
+    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (internal_logfire_code_example.__file__,)
+
+    user_code_example.user1()
+
+    tracebacks = [
+        r.exc_text.replace(  # type: ignore
+            user_code_example.__file__,
+            'user_code_example.py',
+        ).replace(
+            internal_logfire_code_example.__file__,
+            'internal_logfire_code_example.py',
+        )
+        for r in caplog.records
+    ]
+
+    # Important notes about these tracebacks:
+    # - They should look very similar to each other, regardless of how log_internal_error was called.
+    # - They should include all frames from internal_logfire_code_example.py.
+    # - They should include exactly 3 frames from user_code_example.py.
+    # - They should look seamless, with each frame pointing to the next one.
+    # - There should be no sign of logfire's internal error handling code.
+    # - The two files should be isolated and stable so that the exact traceback contents can be asserted.
+    assert tracebacks == snapshot(
+        [
+            """\
+Traceback (most recent call last):
+  File "user_code_example.py", line 22, in user4
+    user5()
+  File "user_code_example.py", line 26, in user5
+    user6()
+  File "user_code_example.py", line 30, in user6
+    outer2(using_decorator)
+  File "internal_logfire_code_example.py", line 36, in outer2
+    outer1(func)
+  File "internal_logfire_code_example.py", line 32, in outer1
+    func()
+  File "internal_logfire_code_example.py", line 16, in using_decorator
+    inner2()
+  File "internal_logfire_code_example.py", line 11, in inner2
+    inner1()
+  File "internal_logfire_code_example.py", line 7, in inner1
+    raise ValueError('inner1')
+ValueError: inner1\
+""",
+            """\
+Traceback (most recent call last):
+  File "user_code_example.py", line 22, in user4
+    user5()
+  File "user_code_example.py", line 26, in user5
+    user6()
+  File "user_code_example.py", line 31, in user6
+    outer2(using_context_manager)
+  File "internal_logfire_code_example.py", line 36, in outer2
+    outer1(func)
+  File "internal_logfire_code_example.py", line 32, in outer1
+    func()
+  File "internal_logfire_code_example.py", line 21, in using_context_manager
+    inner2()
+  File "internal_logfire_code_example.py", line 11, in inner2
+    inner1()
+  File "internal_logfire_code_example.py", line 7, in inner1
+    raise ValueError('inner1')
+ValueError: inner1\
+""",
+            """\
+Traceback (most recent call last):
+  File "user_code_example.py", line 22, in user4
+    user5()
+  File "user_code_example.py", line 26, in user5
+    user6()
+  File "user_code_example.py", line 32, in user6
+    outer2(using_try_except)
+  File "internal_logfire_code_example.py", line 36, in outer2
+    outer1(func)
+  File "internal_logfire_code_example.py", line 32, in outer1
+    func()
+  File "internal_logfire_code_example.py", line 26, in using_try_except
+    inner2()
+  File "internal_logfire_code_example.py", line 11, in inner2
+    inner1()
+  File "internal_logfire_code_example.py", line 7, in inner1
+    raise ValueError('inner1')
+ValueError: inner1\
+""",
+        ]
+    )