fix: Improve BigQuery Agent Analytics plugin reliability and code quality#4491
fix: Improve BigQuery Agent Analytics plugin reliability and code quality#4491caohy1988 wants to merge 1 commit intogoogle:mainfrom
Conversation
|
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. |
|
Response from ADK Triaging Agent Hello @caohy1988, thank you for creating this PR! Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). You can visit https://cla.developers.google.com/ to see your current agreements or to sign a new one. Also, this PR is a bug fix. Could you please associate a GitHub issue with this PR? If there is no existing issue, could you please create one? In addition, could you please provide logs or a screenshot after the fix is applied? This information will help reviewers to review your PR more efficiently. Thanks! |
401fb43 to
d9347d7
Compare
| if TYPE_CHECKING: | ||
| from ..agents.invocation_context import InvocationContext | ||
|
|
||
|
|
There was a problem hiding this comment.
remove unnecessary empty line
d9347d7 to
35d6028
Compare
| state_delta: dict[str, Any], | ||
| invocation_context: InvocationContext, | ||
| event: "Event", | ||
| **kwargs, |
There was a problem hiding this comment.
can you please help me understand why have **kwargs here? ADK doesn't set addtional arguments when running the callback
There was a problem hiding this comment.
Good catch — the **kwargs was added as a forward-compatibility guard in case ADK's callback signatures evolve, but you're right that the framework currently passes only the named parameters declared in BasePlugin. Removed **kwargs from all 12 callback signatures to match BasePlugin exactly. Also fixed on_model_error_callback to include the llm_request parameter that the base class declares.
…lity This CL fixes several bugs in the BigQuery Agent Analytics plugin and refactors the internal data-passing pattern for better type safety and maintainability. - **Stale Loop State Validation:** Use `loop.is_closed()` — a public, reliable API — to detect and clean up stale asyncio loop states in `_batch_processor_prop`, `_get_loop_state`, and `flush`. The previous approach used `asyncio.Queue._loop` which is `None` on Python 3.10+, causing the check to always treat states as stale. - **Quota Project ID Fallback:** Remove the `or project_id` fallback when setting `quota_project_id` on `BigQueryWriteAsyncClient`. This fixes Workload Identity Federation flows where the federated identity lacks `serviceusage.services.use` on the quota project. - **Kwargs Passthrough:** Pass `**kwargs` through to `_log_event` in all callbacks. Previously only model callbacks forwarded them, causing custom attributes (e.g. `customer_id`) to silently drop for agent, tool, run, and error events. - **State Delta Logging:** Replace the dead `on_state_change_callback` (never invoked by the framework) with `on_event_callback`, which is already dispatched by the runner for every event. Remove duplicate `STATE_DELTA` logging from `after_tool_callback`. - **EventData Dataclass:** Replace the `**kwargs`-as-data-bus pattern in `_log_event` with an explicit `EventData` dataclass. This makes the interface self-documenting, catches typos at construction time, and eliminates shared dict mutation across `_resolve_span_ids`, `_extract_latency`, and `_enrich_attributes`. All 12 callback call sites now construct typed `EventData` instances. - **Multi-Subagent Tool Logging Tests:** Add `TestMultiSubagentToolLogging` (6 tests) verifying that tool events are correctly attributed to subagents in multi-turn, multi-agent scenarios. Total tests: 111 (up from 60). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
35d6028 to
2f66ad4
Compare
This CL fixes several bugs in the BigQuery Agent Analytics plugin and refactors the internal data-passing pattern for better type safety and maintainability.
loop.is_closed()— a public, reliable API — to detect and clean up stale asyncio loop states in_batch_processor_prop,_get_loop_state, andflush. The previous approach usedasyncio.Queue._loopwhich isNoneon Python 3.10+, causing the check to always treat states as stale.or project_idfallback when settingquota_project_idonBigQueryWriteAsyncClient. This fixes Workload Identity Federation flows where the federated identity lacksserviceusage.services.useon the quota project.**kwargsthrough to_log_eventin all callbacks. Previously only model callbacks forwarded them, causing custom attributes (e.g.customer_id) to silently drop for agent, tool, run, and error events.on_state_change_callback(never invoked by the framework) withon_event_callback, which is already dispatched by the runner for every event. Remove duplicateSTATE_DELTAlogging fromafter_tool_callback.**kwargs-as-data-bus pattern in_log_eventwith an explicitEventDatadataclass. This makes the interface self-documenting, catches typos at construction time, and eliminates shared dict mutation across_resolve_span_ids,_extract_latency, and_enrich_attributes. All 12 callback call sites now construct typedEventDatainstances.TestMultiSubagentToolLogging(6 tests) verifying that tool events are correctly attributed to subagents in multi-turn, multi-agent scenarios. Total tests: 111 (up from 60).