diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 44107080..bc100b8a 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -1599,6 +1599,16 @@ class BaseLlmFlow(ABC): if response := await self._handle_before_model_callback( invocation_context, llm_request, model_response_event ): + # The model was never called, but the span still has to carry its + # attributes: trace consumers key off the event id attribute and + # drop spans that lack it. + trace_call_llm( + invocation_context, + model_response_event.id, + llm_request, + response, + span, + ) yield response return diff --git a/tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py b/tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py index bc19c458..c75ef84e 100644 --- a/tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py +++ b/tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py @@ -43,11 +43,12 @@ _SPAN_ID_INVALID = 0 class _SpanCapture: - """Stores the span ID and trace ID observed from within a callback.""" + """Stores the span observed from within a callback.""" def __init__(self): self.span_id: int = _SPAN_ID_INVALID self.trace_id: int = 0 + self.span: Optional[trace.Span] = None def capture(self): span = trace.get_current_span() @@ -55,6 +56,7 @@ class _SpanCapture: if ctx and ctx.span_id != _SPAN_ID_INVALID: self.span_id = ctx.span_id self.trace_id = ctx.trace_id + self.span = span class SpanCapturingPlugin(BasePlugin): @@ -220,6 +222,51 @@ def test_short_circuit_before_callback_sees_valid_span(): assert plugin.after_capture.span_id == _SPAN_ID_INVALID +def test_short_circuit_call_llm_span_has_attributes(): + """A short-circuited model call still records the call_llm attributes. + + Trace consumers look a span up by its event id attribute and discard spans + that do not carry one, so an attribute-less span is an invisible span. + """ + plugin = SpanCapturingPlugin() + plugin._short_circuit_before = True + plugin._short_circuit_response = LlmResponse( + content=testing_utils.ModelContent( + [types.Part.from_text(text='short_circuited')] + ) + ) + mock_model = testing_utils.MockModel.create(responses=['unused']) + agent = Agent(name='root_agent', model=mock_model) + runner = testing_utils.InMemoryRunner(agent, plugins=[plugin]) + + events = runner.run('test') + + span = plugin.before_capture.span + assert span is not None, 'no call_llm span was captured' + assert span.name == 'call_llm' + attributes = dict(span.attributes or {}) + + model_event_ids = { + event.id for event in events if event.author == 'root_agent' + } + assert model_event_ids, 'the short-circuit response produced no event' + assert attributes.get('gcp.vertex.agent.event_id') in model_event_ids, ( + 'call_llm span carries no event id on the short-circuit path, so the' + f' trace for that event is unreachable; attributes={attributes}' + ) + for key in ( + 'gen_ai.system', + 'gcp.vertex.agent.invocation_id', + 'gcp.vertex.agent.session_id', + 'gcp.vertex.agent.llm_request', + 'gcp.vertex.agent.llm_response', + ): + assert key in attributes, ( + f'call_llm span is missing {key} on the short-circuit path;' + f' attributes={attributes}' + ) + + # --------------------------------------------------------------------------- # Tests: all three callbacks share same span on error path # ---------------------------------------------------------------------------