Added error correlation headers.

rollbar/__init__.py

@@ -1,4 +1,4 @@
-from __future__ import absolute_import
+from __future__ import absolute_import, annotations
 from __future__ import unicode_literals
 
 import copy
@@ -22,7 +22,8 @@
 import requests
 
 from rollbar.lib import events, filters, dict_merge, transport, defaultJSONEncode
-
+from rollbar.lib.payload import Attribute
+from rollbar.lib.session import get_current_session, set_current_session, parse_session_request_baggage_headers
 
 __version__ = '1.3.0'
 __log_name__ = 'rollbar'
@@ -800,6 +801,7 @@ def _report_exc_info(exc_info, request, extra_data, payload_data, level=None):
     _add_request_data(data, request)
     _add_person_data(data, request)
     _add_lambda_context_data(data)
+    _add_session_data(data)
     data['server'] = _build_server_data()
 
     if payload_data:
@@ -881,6 +883,7 @@ def _report_message(message, level, request, extra_data, payload_data):
     _add_request_data(data, request)
     _add_person_data(data, request)
     _add_lambda_context_data(data)
+    _add_session_data(data)
     data['server'] = _build_server_data()
 
     if payload_data:
@@ -892,6 +895,50 @@ def _report_message(message, level, request, extra_data, payload_data):
     return data['uuid']
 
 
+def _add_session_data(data: dict) -> None:
+    """
+    Adds session data to the payload data if it can be found in the current session or request.
+    """
+    session_data = get_current_session()
+    if session_data:
+        _add_session_attributes(data, session_data)
+        return
+
+    request = _session_data_from_request(data)
+    if request is None:
+        return
+    session_data = parse_session_request_baggage_headers(request.get('headers', None))
+
+    if session_data:
+        _add_session_attributes(data, session_data)
+
+
+def _add_session_attributes(data: dict, session_data: list[Attribute]) -> None:
+    """
+    Adds session attributes to the payload data. This function is careful to not overwrite any existing data in the
+    payload.
+    """
+    if 'attributes' not in data:
+        data['attributes'] = session_data
+        return
+
+    existing_keys = {a['key'] for a in data['attributes']}
+
+    for attribute in session_data:
+        if attribute['key'] not in existing_keys:
+            data['attributes'].append(attribute)
+
+
+def _session_data_from_request(data: dict) -> dict:
+    """
+    Tries to find session data in the request object. Use the request object if provided, otherwise check the data as
+    it may already contain the request object. This is true for some frameworks (e.g. Django).
+    """
+    if data is not None and 'request' in data:
+        return data.get('request', None)
+    return _get_actual_request(_build_request_data(get_request()))
+
+
 def _check_config():
     if not SETTINGS.get('enabled'):
         log.info("pyrollbar: Not reporting because rollbar is disabled.")
@@ -1243,10 +1290,7 @@ def _extract_wsgi_headers(items):
 
 
 def _build_django_request_data(request):
-    try:
-        url = request.get_raw_uri()
-    except AttributeError:
-        url = request.build_absolute_uri()
+    url = request.build_absolute_uri()
 
     request_data = {
         'url': url,

rollbar/contrib/asgi/middleware.py

@@ -1,10 +1,12 @@
 import logging
 import sys
+from typing import Iterable
 
 import rollbar
 from .integration import IntegrationBase, integrate
 from .types import ASGIApp, Receive, Scope, Send
 from rollbar.lib._async import RollbarAsyncError, try_report
+from rollbar.lib.session import set_current_session, reset_current_session
 
 log = logging.getLogger(__name__)
 
@@ -17,6 +19,8 @@ def __init__(self, app: ASGIApp) -> None:
         self.app = app
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope['type'] == 'http':
+            set_current_session(self._format_headers(scope['headers']))
         try:
             await self.app(scope, receive, send)
         except Exception:
@@ -31,3 +35,15 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                     )
                     rollbar.report_exc_info(exc_info)
             raise
+        finally:
+            if scope['type'] == 'http':
+                reset_current_session()
+
+    @staticmethod
+    def _format_headers(headers: Iterable[tuple[bytes, bytes]]) -> dict[str, str]:
+        """
+        Convert list of header tuples to a dictionary with string keys and values.
+
+        Headers are expected to be in the format: [(b'header-name', b'header-value'), ...]
+        """
+        return {key.decode('latin-1'): value.decode('latin-1') for key, value in headers}

rollbar/contrib/django/middleware.py

@@ -87,6 +87,9 @@ def get_payload_data(self, request, exc):
 from django.conf import settings
 from django.http import Http404
 
+from rollbar import set_current_session
+from rollbar.lib.session import reset_current_session
+
 try:
     from django.urls import resolve
 except ImportError:
@@ -258,6 +261,20 @@ def hook(request, data):
                     " Exception was: %r", e
                 )
 
+    def __call__(self, request):
+        headers = {}
+        for k, v in request.META.items():
+            if k.startswith('HTTP_'):
+                header_name = '-'.join(k[len('HTTP_'):].replace('_', ' ').title().split(' '))
+                headers[header_name] = v
+        set_current_session(headers)
+
+        try:
+            response = self.get_response(request)
+            return response
+        finally:
+            reset_current_session()
+
     def _ensure_log_handler(self):
         """
         If there's no log configuration, set up a default handler.

rollbar/contrib/pyramid/__init__.py

@@ -11,6 +11,8 @@
 from pyramid.settings import asbool
 
 import rollbar
+from rollbar import set_current_session
+from rollbar.lib.session import reset_current_session
 
 DEFAULT_WEB_BASE = 'https://rollbar.com'
 BOOLEAN_SETTINGS = [
@@ -50,6 +52,7 @@ def rollbar_tween_factory(pyramid_handler, registry):
     settings = parse_settings(registry.settings)
 
     def rollbar_tween(request):
+        set_current_session(dict(request.headers))
         # for testing out the integration
         try:
             if (settings.get('allow_test', 'true') == 'true' and
@@ -66,6 +69,8 @@ def rollbar_tween(request):
         except Exception as exc:
             handle_error(request, exc, sys.exc_info())
             raise
+        finally:
+            reset_current_session()
         if request.exception is not None:
             handle_error(request, request.exception, request.exc_info)
         return response

rollbar/contrib/starlette/middleware.py

@@ -10,13 +10,16 @@
 from rollbar.contrib.asgi import ReporterMiddleware as ASGIReporterMiddleware
 from rollbar.contrib.asgi.integration import integrate
 from rollbar.lib._async import RollbarAsyncError, try_report
+from rollbar.lib.session import set_current_session, reset_current_session
 
 log = logging.getLogger(__name__)
 
 
 @integrate(framework_name=f'starlette {__version__}')
 class ReporterMiddleware(ASGIReporterMiddleware):
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope['type'] == 'http':
+            set_current_session(self._format_headers(scope['headers']))
         try:
             store_current_request(scope, receive)
             await self.app(scope, receive, send)
@@ -43,3 +46,6 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                     )
                     rollbar.report_exc_info(exc_info, request)
             raise
+        finally:
+            if scope['type'] == 'http':
+                reset_current_session()

rollbar/lib/payload.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Attribute(TypedDict):
+    """
+    Represents the `data.attributes` field in the payload, which is used to store session, execution scope information,
+    and other key-value pairs.
+    """
+    key: str
+    value: str

rollbar/lib/session.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import random
+import threading
+from contextvars import ContextVar
+
+from rollbar.lib.payload import Attribute
+
+_context_session: ContextVar[list[Attribute]|None] = ContextVar('rollbar-session', default=None)
+_thread_session: threading.local = threading.local()
+
+
+def set_current_session(headers: dict[str, str]) -> None:
+    """
+    Set current session data.
+
+    The session data should be a dictionary with string keys and string values.
+    """
+    session_data = parse_session_request_baggage_headers(headers)
+    _context_session.set(session_data)
+    _thread_session.data = session_data
+
+
+def get_current_session() -> list[Attribute]:
+    """
+    Return current session data.
+
+    Do NOT modify the returned session data.
+    """
+    session_data = _context_session.get()
+    if session_data is not None:
+        return session_data
+
+    # Fallback to thread local storage for non-async contexts.
+    return getattr(_thread_session, 'data', None) or []
+
+
+def reset_current_session() -> None:
+    """
+    Reset current session data.
+    """
+    _context_session.set(None)
+    _thread_session.data = None
+
+
+def parse_session_request_baggage_headers(headers: dict) -> list[Attribute]:
+    """
+    Parse the 'baggage' header from the request headers to extract session information. If the 'baggage' header is not
+    present or does not contain the expected keys, a new execution scope ID will be generated and returned as part of
+    the session attributes.
+    """
+    if not headers:
+        return _build_new_scope_attributes()
+
+    baggage_header = None
+
+    # Make sure to handle case-insensitive header keys.
+    for key in headers.keys():
+        if key.lower() == 'baggage':
+            baggage_header = headers[key]
+            break
+
+    if not baggage_header:
+        return _build_new_scope_attributes()
+
+    baggage_items = baggage_header.split(',')
+    baggage_data = []
+    has_scope_id = False
+    for item in baggage_items:
+        if '=' not in item:
+            continue
+        key, value = item.split('=', 1)
+        key = key.strip()
+        if key == 'rollbar.session.id':
+            baggage_data.append({'key': 'session_id', 'value': value.strip()})
+        if key == 'rollbar.execution.scope.id':
+            has_scope_id = True
+            baggage_data.append({'key': 'execution_scope_id', 'value': value.strip()})
+
+    if not baggage_data:
+        return _build_new_scope_attributes()
+
+    # Always ensure we have an execution scope ID, even if the baggage header is present but doesn't contain it.
+    if not has_scope_id:
+        baggage_data.extend(_build_new_scope_attributes())
+
+    return baggage_data
+
+
+def _build_new_scope_attributes() -> list[Attribute]:
+    """
+    Generates a new value for the `rollbar.execution.scope.id` attribute.
+    """
+    new_id = _new_scope_id()
+    if new_id is None:
+        return []
+    return [{'key': 'execution_scope_id', 'value': new_id}]
+
+
+def _new_scope_id() -> str | None:
+    """
+    Generate a new random ID with 128 bits of randomness, formatted as a 32-character hexadecimal string. To be used as
+    an execution scope ID.
+    """
+    try:
+        # Generate a random integer with exactly 128 random bits
+        num = random.getrandbits(128)
+    except Exception as e:
+        return None
+    return format(num, "032x")

rollbar/test/asgi_tests/test_middleware.py

@@ -31,7 +31,7 @@ def test_should_catch_and_report_errors(self, mock_report):
         testapp = ReporterMiddleware(FailingTestASGIApp())
 
         with self.assertRaises(RuntimeError):
-            run(testapp({'type': 'http'}, None, None))
+            run(testapp({'type': 'http', 'headers': []}, None, None))
 
         self.assertTrue(mock_report.called)
 
@@ -73,7 +73,7 @@ def test_should_use_async_report_exc_info_if_default_handler(
         testapp = ReporterMiddleware(FailingTestASGIApp())
 
         with self.assertRaises(RuntimeError):
-            run(testapp({'type': 'http'}, None, None))
+            run(testapp({'type': 'http', 'headers': []}, None, None))
 
         self.assertTrue(async_report_exc_info.called)
         self.assertFalse(sync_report_exc_info.called)
@@ -92,7 +92,7 @@ def test_should_use_async_report_exc_info_if_any_async_handler(
         testapp = ReporterMiddleware(FailingTestASGIApp())
 
         with self.assertRaises(RuntimeError):
-            run(testapp({'type': 'http'}, None, None))
+            run(testapp({'type': 'http', 'headers': []}, None, None))
 
         self.assertTrue(async_report_exc_info.called)
         self.assertFalse(sync_report_exc_info.called)
@@ -112,7 +112,7 @@ def test_should_use_sync_report_exc_info_if_non_async_handlers(
         testapp = ReporterMiddleware(FailingTestASGIApp())
 
         with self.assertRaises(RuntimeError):
-            run(testapp({'type': 'http'}, None, None))
+            run(testapp({'type': 'http', 'headers': []}, None, None))
 
         self.assertFalse(async_report_exc_info.called)
         self.assertTrue(sync_report_exc_info.called)
@@ -128,7 +128,7 @@ def test_should_support_http_only(self):
 
         with mock.patch('rollbar.report_exc_info') as mock_report:
             with self.assertRaises(RuntimeError):
-                run(testapp({'type': 'http'}, None, None))
+                run(testapp({'type': 'http', 'headers': []}, None, None))
 
             self.assertTrue(mock_report.called)
 

rollbar/test/fastapi_tests/test_middleware.py

@@ -370,7 +370,7 @@ def test_should_support_http_only(self):
 
         with mock.patch('rollbar.report_exc_info') as mock_report:
             with self.assertRaises(RuntimeError):
-                run(testapp({'type': 'http'}, None, None))
+                run(testapp({'type': 'http', 'headers': []}, None, None))
 
             mock_report.assert_called_once()
 

rollbar/test/fastapi_tests/test_utils.py

@@ -162,7 +162,6 @@ def test_is_current_version_higher_or_equal(self):
 
         previous_version = None
         for version in versions:
-            print(f'{version} >= {previous_version}')
             if previous_version is None:
                 previous_version = version
                 continue