Class ContextThread

src/memos/context/context.py:206–241

Thread class that automatically propagates the main thread's trace_id to child threads.

class ContextThread(threading.Thread):
    """
    Thread class that automatically propagates the main thread's trace_id to child threads.
    """

    def __init__(self, target, args=(), kwargs=None, **thread_kwargs):
        super().__init__(**thread_kwargs)
        self.target = target
        self.args = args
        self.kwargs = kwargs or {}

        self.main_trace_id = get_current_trace_id()
        self.main_api_path = get_current_api_path()
        self.main_env = get_current_env()
        self.main_user_type = get_current_user_type()
        self.main_user_name = get_current_user_name()
        self.main_context = get_current_context()

    def run(self):
        # Create a new RequestContext with the main thread's trace_id
        if self.main_context:
            # Copy the context data
            child_context = RequestContext(
                trace_id=self.main_trace_id,
                api_path=self.main_api_path,
                env=self.main_env,
                user_type=self.main_user_type,
                user_name=self.main_user_name,
            )
            child_context._data = self.main_context._data.copy()

            # Set the context in the child thread
            set_request_context(child_context)

        # Run the target function
        self.target(*self.args, **self.kwargs)


class ContextThreadPoolExecutor(ThreadPoolExecutor):