[ADK BigQuery Plugin] Event trace_id fractures between INVOCATION_STARTING and AGENT_STARTING

[evekhm]

There is a bug in the bigquery_agent_analytics_plugin where the trace_id for a single invocation splits halfway through the turn.

Early lifecycle events (USER_MESSAGE_RECEIVED and INVOCATION_STARTING) are logged with one
trace_id (fallback to invocation_id), while subsequent events (starting at AGENT_STARTING and continuing through LLM/Tool interactions) are logged under a completely different, newly generated trace_id.

Impact: This breaks downstream joins and observability.
A trace_id is meant to be the single, unbroken thread that ties a user's action to the database logs. By falling back to the invocation_id early, but generating an OTel trace ID late, the plugin breaks the data model in half. You cannot easily GROUP BY trace_id to evaluate an entire turn without writing complex queries holding the invocation_id together.

You can verify the data corruption in BigQuery using the following query on the destination table:

WITH SequencedEvents AS (
  SELECT
    timestamp,
    session_id,
    invocation_id,
    event_type,
    trace_id,
    LEAD(event_type) OVER(PARTITION BY invocation_id ORDER BY timestamp ASC) as next_event,
    LEAD(trace_id) OVER(PARTITION BY invocation_id ORDER BY timestamp ASC) as next_trace_id
  FROM `YOUR_PROJECT.YOUR_DATASET.agent_events`
  WHERE event_type IN ('INVOCATION_STARTING', 'AGENT_STARTING')
)
SELECT *
FROM SequencedEvents
WHERE event_type = 'INVOCATION_STARTING'
  AND next_event = 'AGENT_STARTING'
  AND trace_id != next_trace_id
ORDER BY timestamp DESC

[caohy1988]

Root Cause

When OTel IS configured but no ambient span exists (e.g. Agent Engine, custom runners), early events (USER_MESSAGE_RECEIVED, INVOCATION_STARTING) fell
back to invocation_id as trace_id, while AGENT_STARTING got a new OTel hex trace_id from push_span("agent") — causing trace_id fracture within a single
invocation.

Additionally, push_span() always created independent root spans (each with its own trace_id), so even when a parent span existed on the internal stack,
child spans wouldn't share the same trace_id.

Changes to bigquery_agent_analytics_plugin.py

1. TraceManager.push_span() — Modified to create child spans when a parent exists on the internal stack. Uses trace.set_span_in_context() to pass the
   parent span context without attaching to the ambient OTel context (preserving Span Hierarchy and Duration Inconsistency in BigQueryAgentAnalyticsPlugin #4561 isolation).
2. TraceManager.ensure_invocation_span() — New static method that ensures a root span exists on the plugin stack before any events are logged:
   - If stack already has entries → no-op
   - If ambient OTel span is valid → attach_current_span() (reuse runner's span)
   - Otherwise → push_span("invocation") (create new root span)
3. on_user_message_callback() — Calls ensure_invocation_span() before logging.
4. before_run_callback() — Calls ensure_invocation_span() before logging.
5. after_run_callback() — Pops the invocation root span and uses span_id_override/parent_span_id_override for the INVOCATION_COMPLETED event.

Changes to test file

1. Updated existing test_otel_integration to expect context=None in start_span call.
2. Added TestTraceIdContinuity class with 3 tests:
   - test_trace_id_continuity_no_ambient_span — The core [ADK BigQuery Plugin] Event trace_id fractures between INVOCATION_STARTING and AGENT_STARTING #4645 regression test
   - test_trace_id_continuity_with_ambient_span — Verifies ambient span is reused
   - test_invocation_root_span_isolated_across_turns — Verifies each turn gets a unique trace_id

[caohy1988]

Trace-ID fix details (#4645)

Root cause: When OTel is configured but the Runner's ambient span is absent (Agent Engine, custom runners), early events fell back to invocation_id while push_span("agent") created a new OTel span with a different hex trace_id.

Fix:

• TraceManager.ensure_invocation_span() — called in on_user_message_callback and before_run_callback before any events fire. Attaches the ambient span if valid, or creates a new root span.
• push_span() — now creates child spans under existing stack parent via trace.set_span_in_context(), so all spans share one trace_id.
• _resolve_ids() (replaces _resolve_span_ids()) — 3-layer priority for trace_id/span_id/parent_span_id:
  1. EventData overrides (post-pop callbacks like after_run_callback)
  2. Ambient OTel span (aligns BQ rows with Cloud Trace when o11y is active)
  3. Plugin's internal span stack (fallback for no-ambient paths)
• EventData.trace_id_override — captured before popping the invocation span in after_run_callback, ensuring INVOCATION_COMPLETED shares the same trace_id.

[haiyuan-eng-google]

fixed