fix: bypass is_tracing_enabled_in_context for first-time deferred batch init

First-time users have is_tracing_enabled_in_context() = False by design
  (it's a prerequisite for should_auto_collect_first_time_traces). This
  caused _initialize_backend_batch to return early without creating the
  batch, and _send_events_to_backend to send to a non-existent batch.

  Add skip_context_check parameter to _initialize_backend_batch so the
  first-time handler can bypass the guard during deferred init. Gate
  backend_initialized on trace_batch_id being set. Call
  _finalize_backend_batch directly instead of finalize_batch (which has
  the same context guard). Sync is_current_batch_ephemeral on success
  to prevent endpoint mismatch between batch creation and event send.
This commit is contained in:
Tiago Freire
2026-03-18 17:40:18 -03:00
parent 45c065f7a3
commit 288605c661
3 changed files with 47 additions and 1 deletions

View File

@@ -100,6 +100,7 @@ class FirstTimeTraceHandler:
user_context=user_context,
execution_metadata=execution_metadata,
use_ephemeral=self.batch_manager.is_current_batch_ephemeral,
skip_context_check=True,
)
if not self.batch_manager.trace_batch_id:

View File

@@ -109,10 +109,11 @@ class TraceBatchManager:
user_context: dict[str, str],
execution_metadata: dict[str, Any],
use_ephemeral: bool = False,
skip_context_check: bool = False,
) -> None:
"""Send batch initialization to backend"""
if not is_tracing_enabled_in_context():
if not skip_context_check and not is_tracing_enabled_in_context():
return
if not self.plus_api or not self.current_batch:

View File

@@ -658,6 +658,16 @@ class TestTraceListenerSetup:
trace_listener.first_time_handler.collected_events = True
mock_batch_response = MagicMock()
mock_batch_response.status_code = 201
mock_batch_response.json.return_value = {
"trace_id": "mock-trace-id",
"ephemeral_trace_id": "mock-ephemeral-trace-id",
"access_code": "TRACE-mock",
}
mock_events_response = MagicMock()
mock_events_response.status_code = 200
with (
patch.object(
trace_listener.first_time_handler,
@@ -667,6 +677,40 @@ class TestTraceListenerSetup:
patch.object(
trace_listener.first_time_handler, "_display_ephemeral_trace_link"
) as mock_display_link,
patch.object(
trace_listener.batch_manager.plus_api,
"initialize_trace_batch",
return_value=mock_batch_response,
),
patch.object(
trace_listener.batch_manager.plus_api,
"initialize_ephemeral_trace_batch",
return_value=mock_batch_response,
),
patch.object(
trace_listener.batch_manager.plus_api,
"send_trace_events",
return_value=mock_events_response,
),
patch.object(
trace_listener.batch_manager.plus_api,
"send_ephemeral_trace_events",
return_value=mock_events_response,
),
patch.object(
trace_listener.batch_manager.plus_api,
"finalize_trace_batch",
return_value=mock_events_response,
),
patch.object(
trace_listener.batch_manager.plus_api,
"finalize_ephemeral_trace_batch",
return_value=mock_events_response,
),
patch.object(
trace_listener.batch_manager,
"_cleanup_batch_data",
),
):
crew.kickoff()
wait_for_event_handlers()