fix: pass call_llm span to error callbacks and use real TracerProvider in tests

Address review feedback from caohy1988:

1. Pass the call_llm span explicitly to _run_and_handle_error and
   re-activate it with trace.use_span(parent_span) for error callbacks.
   This ensures on_model_error_callback reliably sees the same span
   context as before_model_callback, defending against async context
   propagation issues across generator yield boundaries.

2. Replace mock.patch-based tracer setup in tests with a real global
   TracerProvider. The original tests masked the production code path
   by mocking the tracer — now they validate the same proxy tracer
   behavior used at runtime.

Fixes #4851

Author: brucearctor

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -307,6 +307,7 @@ async def _run_and_handle_error(
     invocation_context: InvocationContext,
     llm_request: LlmRequest,
     model_response_event: Event,
+    parent_span: Optional[trace.Span] = None,
 ) -> AsyncGenerator[LlmResponse, None]:
   """Wraps an LLM response generator with error callback handling.
 
@@ -320,6 +321,9 @@ async def _run_and_handle_error(
     invocation_context: The invocation context.
     llm_request: The LLM request.
     model_response_event: The model response event.
+    parent_span: Optional parent span (e.g. call_llm) to re-activate for
+      error callbacks, ensuring on_model_error_callback sees the same
+      span context as before_model_callback (issue #4851).
 
   Yields:
     LlmResponse objects from the generator.
@@ -381,11 +385,23 @@ async def _run_on_model_error_callbacks(
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
-    error_response = await _run_on_model_error_callbacks(
-        callback_context=callback_context,
-        llm_request=llm_request,
-        error=model_error,
-    )
+    # Re-activate the parent span (call_llm) so on_model_error_callback
+    # sees the same span_id as before_model_callback (issue #4851).
+    # This is necessary because the inference span has exited and async
+    # context propagation across generator yields can be unreliable.
+    if parent_span is not None:
+      with trace.use_span(parent_span):
+        error_response = await _run_on_model_error_callbacks(
+            callback_context=callback_context,
+            llm_request=llm_request,
+            error=model_error,
+        )
+    else:
+      error_response = await _run_on_model_error_callbacks(
+          callback_context=callback_context,
+          llm_request=llm_request,
+          error=model_error,
+      )
     if error_response is not None:
       yield error_response
     else:
@@ -1153,6 +1169,7 @@ async def _apply_after_model_callback(
                   invocation_context,
                   llm_request,
                   model_response_event,
+                  parent_span=span,
               )
           ) as agen:
             async for llm_response in agen:
@@ -1182,6 +1199,7 @@ async def _apply_after_model_callback(
                   invocation_context,
                   llm_request,
                   model_response_event,
+                  parent_span=span,
               )
           ) as agen:
             async for llm_response in agen:
@@ -1247,13 +1265,15 @@ async def _run_and_handle_error(
       invocation_context: InvocationContext,
       llm_request: LlmRequest,
       model_response_event: Event,
+      parent_span: Optional[trace.Span] = None,
   ) -> AsyncGenerator[LlmResponse, None]:
     async with Aclosing(
         _run_and_handle_error(
             response_generator,
             invocation_context,
             llm_request,
             model_response_event,
+            parent_span=parent_span,
         )
     ) as agen:
       async for response in agen:

tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py

@@ -18,18 +18,18 @@
 after_model_callback, and on_model_error_callback must all execute within
 the same call_llm span so that plugins (e.g. BigQueryAgentAnalyticsPlugin)
 see a consistent span_id for LLM_REQUEST and LLM_RESPONSE events.
+
+These tests set up a real TracerProvider globally — rather than mocking
+the tracer — so that they validate the same code path used in production.
 """
 
 from typing import Optional
-from unittest import mock
 
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.llm_agent import Agent
-from google.adk.flows.llm_flows import base_llm_flow
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from google.adk.plugins.base_plugin import BasePlugin
-from google.adk.telemetry import tracing as adk_tracing
 from google.genai import types
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
@@ -38,10 +38,23 @@
 from ... import testing_utils
 
 
-def _make_real_tracer():
-  """Create a real tracer that produces valid span IDs."""
+@pytest.fixture(autouse=True)
+def _setup_real_tracer_provider():
+  """Set up a real TracerProvider globally for realistic span validation.
+
+  This ensures that all code paths — including the module-level ``tracer``
+  in ``tracing.py`` — produce real spans with valid span IDs, matching
+  production behavior when a TracerProvider is configured.
+
+  Note: ``trace.get_tracer()`` returns a proxy that delegates to the
+  currently set TracerProvider, so setting the provider after import
+  correctly affects all existing tracer references.
+  """
   provider = TracerProvider()
-  return provider.get_tracer('test_tracer')
+  previous_provider = trace.get_tracer_provider()
+  trace.set_tracer_provider(provider)
+  yield
+  trace.set_tracer_provider(previous_provider)
 
 
 class SpanCapturingPlugin(BasePlugin):
@@ -104,20 +117,15 @@ async def test_before_and_after_model_callbacks_share_span_id():
   mismatch between LLM_REQUEST and LLM_RESPONSE events.
   """
   plugin = SpanCapturingPlugin()
-  real_tracer = _make_real_tracer()
 
   mock_model = testing_utils.MockModel.create(responses=['model_response'])
   agent = Agent(
       name='test_agent',
       model=mock_model,
   )
 
-  with (
-      mock.patch.object(base_llm_flow, 'tracer', real_tracer),
-      mock.patch.object(adk_tracing, 'tracer', real_tracer),
-  ):
-    runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
-    events = await runner.run_async_with_new_session('test')
+  runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
+  events = await runner.run_async_with_new_session('test')
 
   # Both callbacks should have captured a span ID
   assert (
@@ -144,7 +152,6 @@ async def test_before_and_on_error_model_callbacks_share_span_id():
   same span as before_model_callback.
   """
   plugin = SpanCapturingPlugin()
-  real_tracer = _make_real_tracer()
 
   mock_model = testing_utils.MockModel.create(
       responses=[], error=SystemError('model error')
@@ -154,12 +161,8 @@ async def test_before_and_on_error_model_callbacks_share_span_id():
       model=mock_model,
   )
 
-  with (
-      mock.patch.object(base_llm_flow, 'tracer', real_tracer),
-      mock.patch.object(adk_tracing, 'tracer', real_tracer),
-  ):
-    runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
-    events = await runner.run_async_with_new_session('test')
+  runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
+  events = await runner.run_async_with_new_session('test')
 
   # Both callbacks should have captured a span ID
   assert (
@@ -206,20 +209,15 @@ async def before_model_callback(
       )
 
   plugin = ShortCircuitPlugin()
-  real_tracer = _make_real_tracer()
 
   mock_model = testing_utils.MockModel.create(responses=['model_response'])
   agent = Agent(
       name='test_agent',
       model=mock_model,
   )
 
-  with (
-      mock.patch.object(base_llm_flow, 'tracer', real_tracer),
-      mock.patch.object(adk_tracing, 'tracer', real_tracer),
-  ):
-    runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
-    events = await runner.run_async_with_new_session('test')
+  runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
+  events = await runner.run_async_with_new_session('test')
 
   # The callback should have a valid (non-zero) span ID from the call_llm span
   assert plugin.span_id is not None and plugin.span_id != 0, (