Class TraceContext

agentops/sdk/core.py:36–72

# Define TraceContext to hold span and token
class TraceContext:
    def __init__(self, span: Span, token: Optional[context_api.Token] = None, is_init_trace: bool = False):
        self.span = span
        self.token = token
        self.is_init_trace = is_init_trace  # Flag to identify the auto-started trace
        self._end_state = StatusCode.UNSET  # Default end state because we don't know yet

    def __enter__(self) -> "TraceContext":
        """Enter the trace context."""
        return self

    def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[Any]) -> bool:
        """Exit the trace context and end the trace.

        Automatically sets the trace status based on whether an exception occurred:
        - If an exception is present, sets status to ERROR
        - If no exception occurred, sets status to OK

        Returns:
            False: Always returns False to propagate any exceptions that occurred
                  within the context manager block, following Python's
                  context manager protocol for proper exception handling.
        """
        if exc_type is not None:
            self._end_state = StatusCode.ERROR
            if exc_val:
                logger.debug(f"Trace exiting with exception: {exc_val}")
        else:
            # No exception occurred, set to OK
            self._end_state = StatusCode.OK

        try:
            tracer.end_trace(self, self._end_state)
        except Exception as e:
            logger.error(f"Error ending trace in context manager: {e}")

        return False


# get_imported_libraries moved to agentops.helpers.system