test: verify callbacks observe the call_llm span by name and ID

Add span name assertions so tests prove each callback sees specifically
the 'call_llm' span (not just any span with a matching ID). This directly
addresses the reviewer's request for proof that before_model_callback,
after_model_callback, and on_model_error_callback all observe the same
valid call_llm span ID.

Author: brucearctor

tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py

@@ -58,24 +58,34 @@ def _setup_real_tracer_provider():
 
 
 class SpanCapturingPlugin(BasePlugin):
-  """Plugin that captures the current span ID in each model callback."""
+  """Plugin that captures span ID and name in each model callback."""
 
   def __init__(self):
     super().__init__(name='span_capturing_plugin')
     self.before_model_span_id: Optional[int] = None
+    self.before_model_span_name: Optional[str] = None
     self.after_model_span_id: Optional[int] = None
+    self.after_model_span_name: Optional[str] = None
     self.on_model_error_span_id: Optional[int] = None
+    self.on_model_error_span_name: Optional[str] = None
+
+  def _capture_span(self):
+    """Capture the current span's ID and name."""
+    span = trace.get_current_span()
+    ctx = span.get_span_context()
+    span_id = ctx.span_id if ctx and ctx.span_id else None
+    span_name = getattr(span, 'name', None)
+    return span_id, span_name
 
   async def before_model_callback(
       self,
       *,
       callback_context: CallbackContext,
       llm_request: LlmRequest,
   ) -> Optional[LlmResponse]:
-    span = trace.get_current_span()
-    ctx = span.get_span_context()
-    if ctx and ctx.span_id:
-      self.before_model_span_id = ctx.span_id
+    self.before_model_span_id, self.before_model_span_name = (
+        self._capture_span()
+    )
     return None
 
   async def after_model_callback(
@@ -84,10 +94,9 @@ async def after_model_callback(
       callback_context: CallbackContext,
       llm_response: LlmResponse,
   ) -> Optional[LlmResponse]:
-    span = trace.get_current_span()
-    ctx = span.get_span_context()
-    if ctx and ctx.span_id:
-      self.after_model_span_id = ctx.span_id
+    self.after_model_span_id, self.after_model_span_name = (
+        self._capture_span()
+    )
     return None
 
   async def on_model_error_callback(
@@ -97,10 +106,9 @@ async def on_model_error_callback(
       llm_request: LlmRequest,
       error: Exception,
   ) -> Optional[LlmResponse]:
-    span = trace.get_current_span()
-    ctx = span.get_span_context()
-    if ctx and ctx.span_id:
-      self.on_model_error_span_id = ctx.span_id
+    self.on_model_error_span_id, self.on_model_error_span_name = (
+        self._capture_span()
+    )
     return LlmResponse(
         content=testing_utils.ModelContent(
             [types.Part.from_text(text='error handled')]
@@ -110,7 +118,7 @@ async def on_model_error_callback(
 
 @pytest.mark.asyncio
 async def test_before_and_after_model_callbacks_share_span_id():
-  """Verify before_model_callback and after_model_callback share the same span.
+  """Verify before_model_callback and after_model_callback share call_llm span.
 
   This is the core regression test for issue #4851. Before the fix,
   before_model_callback ran outside the call_llm span, causing a span_id
@@ -125,16 +133,26 @@ async def test_before_and_after_model_callbacks_share_span_id():
   )
 
   runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
-  events = await runner.run_async_with_new_session('test')
+  await runner.run_async_with_new_session('test')
 
-  # Both callbacks should have captured a span ID
+  # Both callbacks must have captured a valid span ID
   assert (
       plugin.before_model_span_id is not None
   ), 'before_model_callback did not capture a span ID'
   assert (
       plugin.after_model_span_id is not None
   ), 'after_model_callback did not capture a span ID'
 
+  # Both must observe the call_llm span specifically
+  assert plugin.before_model_span_name == 'call_llm', (
+      f'before_model_callback saw span "{plugin.before_model_span_name}", '
+      'expected "call_llm"'
+  )
+  assert plugin.after_model_span_name == 'call_llm', (
+      f'after_model_callback saw span "{plugin.after_model_span_name}", '
+      'expected "call_llm"'
+  )
+
   # The span IDs must match — this is the core assertion for issue #4851
   assert plugin.before_model_span_id == plugin.after_model_span_id, (
       'Span ID mismatch: before_model_callback span_id='
@@ -146,10 +164,10 @@ async def test_before_and_after_model_callbacks_share_span_id():
 
 @pytest.mark.asyncio
 async def test_before_and_on_error_model_callbacks_share_span_id():
-  """Verify before_model_callback and on_model_error_callback share span.
+  """Verify before_model_callback and on_model_error_callback share call_llm span.
 
   When the model raises an error, on_model_error_callback should see the
-  same span as before_model_callback.
+  same call_llm span as before_model_callback.
   """
   plugin = SpanCapturingPlugin()
 
@@ -162,16 +180,26 @@ async def test_before_and_on_error_model_callbacks_share_span_id():
   )
 
   runner = testing_utils.TestInMemoryRunner(agent, plugins=[plugin])
-  events = await runner.run_async_with_new_session('test')
+  await runner.run_async_with_new_session('test')
 
-  # Both callbacks should have captured a span ID
+  # Both callbacks must have captured a valid span ID
   assert (
       plugin.before_model_span_id is not None
   ), 'before_model_callback did not capture a span ID'
   assert (
       plugin.on_model_error_span_id is not None
   ), 'on_model_error_callback did not capture a span ID'
 
+  # Both must observe the call_llm span specifically
+  assert plugin.before_model_span_name == 'call_llm', (
+      f'before_model_callback saw span "{plugin.before_model_span_name}", '
+      'expected "call_llm"'
+  )
+  assert plugin.on_model_error_span_name == 'call_llm', (
+      f'on_model_error_callback saw span "{plugin.on_model_error_span_name}", '
+      'expected "call_llm"'
+  )
+
   # The span IDs must match
   assert plugin.before_model_span_id == plugin.on_model_error_span_id, (
       'Span ID mismatch: before_model_callback span_id='
@@ -184,13 +212,14 @@ async def test_before_and_on_error_model_callbacks_share_span_id():
 
 @pytest.mark.asyncio
 async def test_before_model_callback_short_circuit_has_span():
-  """Verify before_model_callback has a valid span when short-circuiting."""
+  """Verify before_model_callback has a valid call_llm span when short-circuiting."""
 
   class ShortCircuitPlugin(BasePlugin):
 
     def __init__(self):
       super().__init__(name='short_circuit_plugin')
       self.span_id: Optional[int] = None
+      self.span_name: Optional[str] = None
 
     async def before_model_callback(
         self,
@@ -202,6 +231,7 @@ async def before_model_callback(
       ctx = span.get_span_context()
       if ctx and ctx.span_id:
         self.span_id = ctx.span_id
+      self.span_name = getattr(span, 'name', None)
       return LlmResponse(
           content=testing_utils.ModelContent(
               [types.Part.from_text(text='short-circuited')]
@@ -225,6 +255,12 @@ async def before_model_callback(
       'short-circuiting the LLM call'
   )
 
+  # Must be specifically the call_llm span
+  assert plugin.span_name == 'call_llm', (
+      f'before_model_callback saw span "{plugin.span_name}", '
+      'expected "call_llm"'
+  )
+
   # Verify the short-circuit response was received
   simplified = testing_utils.simplify_events(events)
   assert any('short-circuited' in str(e) for e in simplified)