fix(telemetry): ensure all model callbacks observe the same call_llm span

[caohy1988]

Summary

• Moves before_model_callback inside the call_llm span so plugins see a valid span (previously ran before the span started)
• Wraps after_model_callback with trace.use_span(span, end_on_exit=False) to rebind to the call_llm span (previously could see a child inference span)
• Threads call_llm_span into _run_and_handle_error and wraps on_model_error_callback the same way
• Adds 5 regression tests proving all three callbacks observe the same call_llm span ID

Fixes #4851

Changes

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

• _call_llm_async(): Moved before_model_callback, label setup, and __get_llm() inside the with tracer.start_as_current_span('call_llm') as span: block
• Both CFC and non-CFC after_model_callback call sites wrapped with trace.use_span(span, end_on_exit=False)
• _run_and_handle_error(): New optional call_llm_span parameter; error callbacks run under use_span(call_llm_span) when provided
• Class method wrapper updated to pass through call_llm_span

tests/unittests/flows/llm_flows/test_llm_callback_span_consistency.py (new):

1. test_before_and_after_callbacks_share_same_span — success path
2. test_callbacks_same_trace_id — same trace
3. test_before_and_error_callbacks_share_same_span — error path
4. test_short_circuit_before_callback_sees_valid_span — short-circuit path
5. test_all_three_callbacks_share_span_on_error — error recovery path (all 3 callbacks)

Design decisions

• end_on_exit=False: The call_llm span is already managed by the with tracer.start_as_current_span(...) context manager. use_span is only used to re-activate the span context, not to end it.
• Backward-compatible: call_llm_span defaults to None; when absent, behavior is unchanged.
• Plugin workaround preserved: This PR does NOT remove the ambient-span workaround in BigQueryAgentAnalyticsPlugin. That should be a separate follow-up PR after this fix is proven in production.

Test plan

• All 5 new span consistency tests pass
• All 363 existing llm_flows tests pass (0 regressions)
• Autoformat clean

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[adk-bot]

Response from ADK Triaging Agent

Hello @caohy1988, thank you for creating this PR!

It looks like you haven't signed a Contributor License Agreement (CLA) yet. You'll need to sign one before we can review your pull request. Please visit https://cla.developers.google.com/ to sign a new one.

Thanks!