Skip to content

Commit 80c5a24

Browse files
haiyuan-eng-googlecopybara-github
authored andcommitted
feat: Enhance BQ plugin with fork safety, auto views, and trace continuity
- **Fork-safety (#4636):** Adds PID tracking to `BigQueryAgentAnalyticsPlugin` so that forked child processes detect stale gRPC channels and re-initialize instead of deadlocking. Uses the standard `os.getpid()` pattern (same as SQLAlchemy, gRPC-python). - **Auto-create views (#4639):** After ensuring the `agent_events` table exists, automatically creates 15 per-event-type BigQuery views (`v_llm_request`, `v_tool_completed`, etc.) that unnest JSON columns into typed, queryable columns. Controlled by `BigQueryLoggerConfig.create_views` (default `True`), idempotent via `CREATE OR REPLACE VIEW`. - **Trace-ID continuity & o11y alignment (#4645):** Fixes trace_id fracture between early events (USER_MESSAGE_RECEIVED, INVOCATION_STARTING) and later events (AGENT_STARTING onwards) when no ambient OTel span exists. Also aligns BQ rows with Cloud Trace span IDs when o11y is active. - **Span-ID consistency under ambient OTel (#4640 review):** Fixes `*_STARTING` / `*_COMPLETED` events producing mismatched span IDs when an ambient OTel span is active. Completion callbacks now check for ambient spans and defer to `_resolve_ids` Layer 2 instead of overriding with plugin-synthetic IDs. - **Stack leak safety (#4640 review):** Adds `TraceManager.clear_stack()` and makes `ensure_invocation_span()` clear stale records from *different* invocations, preventing span stack leaks across invocations. Uses `_active_invocation_id_ctx` to distinguish stale leak vs same-invocation re-entry. - **Root agent name staleness fix:** `init_trace()` now refreshes `_root_agent_name_ctx` unconditionally on each invocation (previously set-once-on-None). `after_run_callback` resets it alongside other invocation cleanup. - **Exception-safe cleanup:** `after_run_callback` uses `try/finally` to guarantee invocation state (`clear_stack`, `_active_invocation_id_ctx`, `_root_agent_name_ctx`) is always reset, even if `_log_event` raises. - **`on_tool_error_callback` span fix:** Previously discarded the span_id from `pop_span()`, causing TOOL_ERROR events to get the wrong span. Now captures and uses the popped span_id. Co-authored-by: Haiyuan Cao <haiyuan@google.com> PiperOrigin-RevId: 877580395
1 parent c59afc2 commit 80c5a24

File tree

2 files changed

+1853
-105
lines changed

2 files changed

+1853
-105
lines changed

0 commit comments

Comments
 (0)