Preserve RegistrationContext in telemetry worker threads

Merged
MA
2mo ago
+116-4

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Description

This PR fixes an issue where telemetry event collection in worker threads would lose the submitting thread's RegistrationContext, causing the worker to create a throwaway context and re-import rxconfig.py unnecessarily.

The Problem:
When telemetry jobs are submitted via _submit(), they run on a background executor thread that has no RegistrationContext attached. During event collection (e.g., get_bun_path()), a config lookup would trigger ensure_context(), which attaches a new context to the worker thread and re-imports rxconfig.py. This is wasteful since the app's config is already loaded in the submitting thread's context.

The Solution:

  • Capture the submitting thread's RegistrationContext (if any) at submission time via new _current_registration_context() helper
  • Pass this context to _run_suppressed() in the worker thread
  • Attach the context for the duration of the job execution, then detach it to keep the worker thread clean for subsequent jobs
  • This allows event collection to reuse the already-loaded config instead of re-importing

Changes:

  • reflex/utils/telemetry.py:
    • Added _current_registration_context() to safely retrieve the caller's context without creating one
    • Modified _run_suppressed() to accept and manage the registration context around job execution
    • Updated _submit() to capture and forward the context
  • tests/units/test_telemetry.py: Added two new tests verifying context preservation and cleanup
  • tests/units/test_config.py: Updated existing test to avoid counting unrelated background thread loads

Testing

  • Added test_submit_runs_job_in_callers_registration_context() to verify the submitting thread's context is available during job execution
  • Added test_submit_leaves_worker_context_clean_between_jobs() to verify the context is properly detached after each job
  • Updated test_get_config_loads_once_for_shared_context() to exclude unrelated background threads from load counting
  • All existing telemetry tests continue to pass

https://claude.ai/code/session_01Ex7f8t8niToPKPyqXh4Eeq