GoogleCloudPlatform / GoogleCloudPlatform/BigQuery-Agent-Analytics-SDK
Track span-level correlation between BQAA rows and Cloud Trace spans
- Dominant language
- Python
- Stars
- 47
- Forks
- 21
- Avg merge
- 2d 13h
- Merged PRs (30d)
- 33
Description
## Problem
BQAA rows currently correlate with Cloud Trace primarily at the trace level via `trace_id`. That is enough to jump from BigQuery Agent Analytics telemetry to the Cloud Trace trace, but it does not provide a reliable span-level join for most child BQAA rows.
Current behavior in the ADK Python BQAA plugin:
- `trace_id` is inherited from the active ambient OpenTelemetry span when present, so all BQAA rows for an invocation can join to the Cloud Trace trace.
- The root invocation BQAA `span_id` may reuse the ambient OTel span id when an active span exists.
- Child BQAA `span_id` / `parent_span_id` values are internal BQAA correlation IDs used to reconstruct the agent/LLM/tool tree in BigQuery.
- The plugin intentionally does not call `tracer.start_span(...)` or export its own OTel spans, to avoid duplicate/confusing spans in Cloud Trace.
- **Schema-description drift (verified, ADK `main` @`c007a874`):** the `span_id` / `parent_span_id` column *descriptions* in `_get_events_schema()` still read "OpenTelemetry span ID", but `_resolve_ids()` populates them from the plugin's internal span stack and only falls back to the ambient OTel id on the empty-stack/invocation-attach path (`attach_current_span()`). The descriptions should be corrected to say "BQAA internal execution-tree id" as part of this work.
This means the current intended model is:
```text
Cloud Trace <-> BQAA: trace_id
BQAA internal hierarchy: span_id + parent_span_id
```
But users may reasonably expect to join Cloud Trace telemetry to BQAA telemetry at span granularity, especially when investigating latency or service-level behavior around a specific agent/tool/LLM operation.
## Tracking question
Can the SDK and/or producer schema support span-level Cloud Trace correlation without replacing BQAA's internal `span_id` semantics or reintroducing exported duplicate OTel spans?
## Proposed direction
Prefer adding explicit OTel correlation metadata rather than overloading the existing BQAA `span_id`:
```text
trace_id -- existing Cloud Trace trace id
span_id -- existing BQAA internal operation id
parent_span_id -- existing BQAA internal parent operation id
otel_span_id -- ambient OTel span id, if valid at event emission time
otel_parent_span_id -- optional, if available/derivable
```
The SDK could then expose helpers/views for:
```text
Cloud Trace span.span_id == BQAA agent_events.otel_span_id -- best-effort, not a foreign key
```
while preserving:
```text
BQAA execution tree == span_id + parent_span_id
```
A less invasive version could store the OTel fields under `attributes.otel.*` (`attributes.otel.span_id`, `attributes.otel.trace_id`) first, then promote to typed SDK/view columns if the shape proves useful.
**Refinements (validated against ADK `main` @`c007a874`):**
1. **Add `otel_span_id` as a *separate* field — never reinterpret `span_id`.** `span_id` / `parent_span_id` stay the BQAA-internal execution tree.
2. **Capture from `trace.get_current_span().get_span_context()` at row-emission time, only when `is_valid`.** This mirrors the existing `attach_current_span()` gate (which already checks `is_valid`), now generalized beyond the root path.
3. **Per-row-type capture semantics** (each row records the ambient exported span it ran under):
- root/invocation rows can legitimately have `span_id == otel_span_id` (the attach case);
- LLM rows capture the ambient `call_llm` span id (ADK wraps model calls in a real `call_llm` span and rebinds after/error callbacks to it);
- **ordinary** tool rows (`TOOL_STARTING` / `TOOL_COMPLETED`) capture each tool's **own** `execute_tool {tool}` span id — even for parallel calls, since ADK wraps every tool invocation in its own `record_tool_execution()` span (`_instrumentation.py`) inside the `asyncio.gather(...)`, with `before_tool_callback` / `after_tool_callback` running inside that span; the `execute_tool (merged)` span (`functions.py`) is created only *after* the gather to trace ADK's merged response event, and is **not** the per-tool span;
- **resume-side long-running** `TOOL_COMPLETED` rows are a documented exception: a non-HITL `FunctionResponse` arriving in a later user message is emitted as `TOOL_COMPLETED` from `on_user_message_callback`, which does **not** run inside the original `execute_tool {tool}` span. Its row-emission-time `otel_span_id` would therefore capture the **current invocation / user-message** ambient span, not the original tool span — i.e. best-effort/current-invocation, unless the implementation explicitly stores the original OTel tool span id with the pause metadata (the same plugin-bookkeeping caveat as `otel_parent_span_id`);
- no plugin-owned OTel span is started or exported (preserves the #94 no-duplicate-span guarantee).
4. **`otel_parent_span_id` is optional / phase 2.** The OTel API does **not** expose a parent id from a `SpanContext`; `_resolve_ids()` only gets a parent via the span *object* (`getattr(span, "parent", None)`). Deriving `otel_parent_span_id` therefore needs span-object access or plugin-side bookkeeping — defer it.
5. **`otel_span_id` is a best-effort Cloud Trace join key, not a foreign key.** A valid span context can be unsampled and therefore absent from the Cloud Trace export; document this so consumers `LEFT JOIN`.
## Acceptance criteria
- Document the current trace-level join behavior clearly, and correct the stale `span_id` / `parent_span_id` schema descriptions to "BQAA internal execution-tree id".
- Add an explicit `otel_span_id` correlation field (staged: `attributes.otel.*` first, typed column/view later), captured from the ambient span context only when `is_valid`.
- Keep existing `span_id` / `parent_span_id` semantics unchanged for BQAA's internal execution tree.
- Do not require the BQAA plugin to create/export its own OTel spans (preserve the #94 no-duplicate-span guarantee).
- Treat `otel_parent_span_id` as optional/phase 2 (not derivable from `SpanContext` alone).
- Document `otel_span_id` as a best-effort (unsampled spans absent) Cloud Trace join key, with a backward-compatible query pattern (`JSON_VALUE(attributes, '$.otel.span_id')`).
- Tests:
- root invocation rows can have `span_id == otel_span_id`;
- LLM rows capture the ambient `call_llm` span id separately from the BQAA-internal `span_id`;
- tool rows capture each tool's individual `execute_tool {tool}` span id; a **parallel-tools test asserts distinct `otel_span_id`s per tool row** (the `execute_tool (merged)` span is only for the merged response event);
- a **long-running tool completion test** pins the resume-side semantics: a `TOOL_COMPLETED` emitted from a later user message captures the current-invocation ambient span (or the stored original tool span id, if the implementation chooses to persist it with the pause metadata) — not the original tool span by default;
- no plugin-owned OTel span is started or exported.
## Notes
This came up while reviewing ADK docs for the BQAA plugin: the docs now clarify that BQAA child `span_id` values are internal IDs and generally do not align one-to-one with Cloud Trace child span ids. The next product question is whether the SDK should provide a first-class span-level correlation path.
Contributor guide
Assessment
This issue has not been assessed yet.